Quieter scripts, handle multiple extensions, other cleanup for release

flow/compute_flow_sequences.py

@@ -14,11 +14,12 @@
 import numpy as np
 from natsort import natsorted
 from PIL import Image
+from script_utils.common import common_setup
 from tqdm import tqdm
 
 from flow.convert_flo_png import convert_flo, convert_flo_pavel_matlab
 from utils.log import setup_logging
-from utils.misc import IMG_EXTENSIONS
+from utils.misc import glob_ext, IMG_EXTENSIONS
 
 
 @contextlib.contextmanager
@@ -186,7 +187,7 @@ def compute_flownet2_flow(flo_input_outputs, logger, gpu,
 
 
 def compute_sequence_flow(input_dir, output_dir, flow_fn, flow_args, gpu_queue,
-                          logger_name, convert_png, remove_flo, extension):
+                          logger_name, convert_png, remove_flo, extensions):
     """
     Args:
         convert_png (str or False): How to convert .flo files to .png. Choices
@@ -204,8 +205,8 @@ def compute_sequence_flow(input_dir, output_dir, flow_fn, flow_args, gpu_queue,
         assert not remove_flo, (
             'remove_flo is only allowed if convert_png is set')
 
-    image_paths = natsorted(
-        list(input_dir.glob('*' + extension)), key=lambda x: x.stem)
+    image_paths = natsorted(glob_ext(input_dir, extensions, recursive=False),
+                            key=lambda x: x.stem)
     times = {}
     times['start'] = time.time()
     file_logger = logging.getLogger(logger_name)
@@ -321,6 +322,7 @@ def main():
               'CPU operations, like loading input/output lists, checking '
               'image dimensions, and converting .flo to .png while other '
               'workers use the GPU.'))
+    parser.add_argument('--quiet', action='store_true')
 
     flownet2_parser = parser.add_argument_group('Flownet2 params')
     flownet2_parser.add_argument(
@@ -347,26 +349,17 @@ def main():
     output_root.mkdir(parents=True, exist_ok=True)
 
     file_name = Path(__file__).stem
-    logging_path = str(
-        output_root /
-        (file_name + '.py.%s.log' % datetime.now().strftime('%b%d-%H-%M-%S')))
-    setup_logging(logging_path)
-    logging.info('Args:\n%s', vars(args))
-
-    args.extensions = [
-        x if x[0] == '.' else ('.' + x) for x in args.extensions
-    ]
-
-    def is_valid(path):
-        return any(path.name.endswith(y) for y in args.extensions)
+    log_level = logging.ERROR if args.quiet else logging.INFO
+    file_logger = common_setup(__file__,
+                               output_root,
+                               args,
+                               log_console_level=log_level)
 
     if args.recursive:
-        sequences = set(x.parent for x in input_root.rglob('*') if is_valid(x))
-        # Handle one-level of symlinks for ease of use.
-        for symlink_dir in input_root.iterdir():
-            if symlink_dir.is_symlink() and symlink_dir.is_dir():
-                sequences.update(x.parent for x in symlink_dir.rglob('*')
-                                 if is_valid(x))
+        sequences = sorted(
+            x.parent
+            for x in glob_ext(input_root, args.extensions, recursive=True)
+        )
     else:
         sequences = sorted(input_root.iterdir())
 
@@ -413,10 +406,10 @@ def is_valid(path):
             'flow_fn': flow_fn,
             'flow_args': flow_args,
             'gpu_queue': gpu_queue,
-            'logger_name': logging_path,
+            'logger_name': file_logger.name,
             'convert_png': convert_png,
             'remove_flo': bool(convert_png),
-            'extension': args.extension
+            'extensions': args.extensions
         })
 
     list(

release/custom/compute_flow.py

@@ -1,7 +1,6 @@
 import argparse
 import logging
 import pprint
-import subprocess
 import yaml
 from pathlib import Path
 
@@ -27,7 +26,8 @@ def main():
         config = yaml.load(f)
 
     args.output_dir.mkdir(exist_ok=True, parents=True)
-    common_setup(__file__, args.output_dir, args)
+    common_setup(__file__, args.output_dir)
+    logging.debug('Args:\n%s', pprint.pformat(vars(args)))
     logging.debug('Config:\n%s', pprint.pformat(config))
 
     msg(f"Computing flow on {args.frames_dir}.")

release/custom/infer.py

@@ -1,10 +1,12 @@
 import argparse
+import logging
+import pprint
 import yaml
 from pathlib import Path
 
 from script_utils.common import common_setup
 
-from release.helpers.misc import subprocess_call
+from release.helpers.misc import msg, subprocess_call
 
 
 def get_config_ckpt(model_dir, step):
@@ -47,7 +49,8 @@ def main():
 
     args = parser.parse_args()
     args.output_dir.mkdir(exist_ok=True, parents=True)
-    common_setup(__file__, args.output_dir, args)
+    common_setup(__file__, args.output_dir)
+    logging.debug('Args:\n%s', pprint.pformat(vars(args)))
 
     detectron_dir = (
         Path(__file__).resolve().parent.parent.parent / 'detectron_pytorch')
@@ -76,6 +79,8 @@ def main():
         image_dirs = ['--image_dirs', args.flow_dir]
         input_type = ['--input_type', 'flow']
 
+    msg(f'Running {args.model} model on {image_dirs[1].resolve()}')
+
     cmd = ['python', 'tools/infer_simple.py']
     args = ([
         '--cfg', model_config,
@@ -86,6 +91,7 @@ def main():
         + input_type + [
         '--save_images', args.visualize,
         '--output_dir', args.output_dir,
+        '--quiet',
         '--recursive'])
     subprocess_call(cmd + args, cwd=str(detectron_dir))
 

release/custom/track.py

@@ -1,4 +1,6 @@
 import argparse
+import logging
+import pprint
 import yaml
 from pathlib import Path
 
@@ -41,7 +43,8 @@ def main():
 
     args.output_dir.mkdir(exist_ok=True, parents=True)
 
-    common_setup(__file__, args.output_dir, args)
+    common_setup(__file__, args.output_dir)
+    logging.debug('Args:\n%s', pprint.pformat(vars(args)))
 
     if args.model == 'joint':
         detections_args = ['--init-detections-dir', args.detections_dir]
@@ -62,7 +65,8 @@ def main():
         '--bidirectional',
         '--score-init-min', 0.9,
         '--fps', args.fps,
-        '--filename-format', args.filename_format
+        '--filename-format', args.filename_format,
+        '--quiet'
         ] + extra_args
     cmd = ['python', script] + args
     msg(f'Running tracker')

release/davis16/evaluate.py

@@ -40,7 +40,9 @@ def main():
     parser = argparse.ArgumentParser(
         description=__doc__.split('\n')[0] if __doc__ else '',
         formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-    parser.add_argument('output_stage', choices=['detection', 'tracking'])
+    parser.add_argument('output_stage',
+                        choices=['detection', 'tracking'],
+                        default='detection')
     parser.add_argument('--config', default=Path('./release/config.yaml'))
 
     args = parser.parse_args()

release/davis16/infer.py

@@ -57,6 +57,7 @@ def main():
             '--input_type', 'rgb', 'flow',
             '--save_images', False,
             '--output_dir', output_dir / split,
+            '--quiet',
             '--recursive'
         ]
         subprocess_call(cmd + args, cwd=str(detectron_dir))

release/davis17/infer.py

@@ -57,6 +57,7 @@ def main():
             '--input_type', 'rgb', 'flow',
             '--save_images', False,
             '--output_dir', output_dir / split,
+            '--quiet',
             '--recursive'
         ]
         subprocess_call(cmd + args, cwd=str(detectron_dir))

release/fbms/infer.py

@@ -55,6 +55,7 @@ def main():
             '--input_type', 'rgb', 'flow',
             '--save_images', False,
             '--output_dir', output_dir / split,
+            '--quiet',
             '--recursive'
         ]
         subprocess_call(cmd + args, cwd=str(detectron_dir))

release/helpers/flow.py

@@ -10,11 +10,11 @@ def compute_flow_helper(config,
                         recursive=True):
     flownet2_dir = Path(config['flow']['flownet2_dir'])
 
-    if isinstance(extensions, (list, tuple)):
-        extensions = ' '.join(extensions)
+    if isinstance(extensions, str):
+        extensions = [extensions]
 
     maybe_recursive = ['--recursive'] if recursive else []
-    maybe_extensions = ['--extensions', extensions] if extensions else []
+    maybe_extensions = (['--extensions'] + extensions) if extensions else []
     args = [
         '--input-dir', input_dir
         ] + maybe_recursive + [
@@ -25,7 +25,8 @@ def compute_flow_helper(config,
         '--output-dir', output_dir,
         '--flow-type', 'flownet2',
         '--flownet2-dir', flownet2_dir,
-        '--flownet2-model', 'kitti'
+        '--flownet2-model', 'kitti',
+        '--quiet'
     ]
     cmd = ['python', 'flow/compute_flow_sequences.py'] + args
     subprocess_call(cmd)

release/helpers/misc.py

@@ -17,9 +17,9 @@ def subprocess_call(cmd, log=True, **kwargs):
         for x in cmd
     ]
     if log:
-        logging.info('Command:\n%s', ' '.join(cmd).replace("--", "\\\n--"))
+        logging.debug('Command:\n%s', ' '.join(cmd).replace("--", "\\\n--"))
         if kwargs:
-            logging.info('subprocess.check_call kwargs:\n%s', kwargs)
+            logging.debug('subprocess.check_call kwargs:\n%s', kwargs)
     subprocess.check_call(cmd, **kwargs)
 
 

release/ytvos/infer.py

@@ -22,11 +22,12 @@ def main():
         formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument('--config', default=Path('./release/config.yaml'))
     parser.add_argument(
-        '--with-ytvos-train',
+        '--without-ytvos-train',
         action='store_true',
-        help=("By default, we infer with model that was not trained on YTVOS "
-              "for fair evaluation of generalization. If --with-ytvos-train "
-              "is specified, infer with model that uses YTVOS for training."))
+        help=("By default, we infer with model that was trained on YTVOS. "
+              "For fair evaluation of generalization, we use a model without "
+              "YTVOS training in our manuscript. Set this to True to use "
+              "the model without YTVOS training."))
 
     args = parser.parse_args()
 
@@ -35,10 +36,10 @@ def main():
 
     args = parser.parse_args()
     output_dir = Path(config['ytvos']['output_dir']) / 'detections'
-    if args.with_ytvos_train:
-        output_dir = output_dir / 'with_ytvos'
-    else:
+    if args.without_ytvos_train:
         output_dir = output_dir / 'without_ytvos'
+    else:
+        output_dir = output_dir / 'with_ytvos'
     output_dir.mkdir(exist_ok=True, parents=True)
     common_setup(__file__, output_dir, args)
 
@@ -73,6 +74,7 @@ def main():
             '--input_type', 'rgb', 'flow',
             '--save_images', False,
             '--output_dir', output_dir / split,
+            '--quiet',
             '--recursive'
         ]
         subprocess_call(cmd + args, cwd=str(detectron_dir))

release/ytvos/track.py

@@ -16,23 +16,23 @@ def main():
         formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument('--config', default=Path('./release/config.yaml'))
     parser.add_argument(
-        '--with-ytvos-train',
+        '--without-ytvos-train',
         action='store_true',
-        help=("By default, we run tracking on outputs from a model that was "
-              "not trained on YTVOS for fair evaluation of generalization. If "
-              "--with-ytvos-train is specified, infer with model that uses "
-              "YTVOS for training."))
+        help=("By default, we infer with model that was trained on YTVOS. "
+              "For fair evaluation of generalization, we use a model without "
+              "YTVOS training in our manuscript. Set this to True to use "
+              "the model without YTVOS training."))
 
     args = parser.parse_args()
 
     with open(args.config, 'r') as f:
         config = yaml.load(f)
 
     output_dir = Path(config['ytvos']['output_dir']) / 'tracks'
-    if args.with_ytvos_train:
-        output_dir = output_dir / 'with_ytvos'
-    else:
+    if args.without_ytvos_train:
         output_dir = output_dir / 'without_ytvos'
+    else:
+        output_dir = output_dir / 'with_ytvos'
     output_dir.mkdir(exist_ok=True, parents=True)
 
     common_setup(__file__, output_dir, args)

tracker/track.py

@@ -24,6 +24,7 @@
 from utils.datasets import get_classes
 from utils.distance import chi_square_distance  # , intersection_distance
 from utils.log import setup_logging
+from utils.misc import IMG_EXTENSIONS
 
 
 class Detection():
@@ -984,7 +985,7 @@ def track_and_visualize(detection_results,
                         images_dir,
                         tracking_params,
                         get_framenumber,
-                        frame_extension,
+                        frame_extensions,
                         vis_dataset=None,
                         output_numpy=None,
                         output_numpy_every_kth=1,
@@ -1007,8 +1008,14 @@ def track_and_visualize(detection_results,
 
     label_list = get_classes(vis_dataset)
 
+    for extension in frame_extensions:
+        if (images_dir / (frames[0] + extension)).exists():
+            break
+    else:
+        raise ValueError(f"Could not find image with any extension at "
+                         f"{images_dir / frames[0]}")
     frame_paths = [
-        images_dir / (frame + frame_extension) for frame in frames
+        images_dir / (frame + extension) for frame in frames
     ]
     all_tracks = track(frame_paths,
                        [detection_results[frame] for frame in frames],
@@ -1056,7 +1063,7 @@ def main():
         '--output-track-file',
         type=Path,
         help='Optional; path to output MOT17 style tracking output.')
-    parser.add_argument('--extension', default='.png')
+    parser.add_argument('--extensions', nargs='*', default=IMG_EXTENSIONS)
     parser.add_argument(
         '--dataset', default='coco', choices=['coco', 'objectness'])
     parser.add_argument(
@@ -1127,7 +1134,7 @@ def get_framenumber(x):
                         args.images_dir,
                         tracking_params,
                         get_framenumber,
-                        args.extension,
+                        args.extensions,
                         vis_dataset=args.dataset,
                         output_images_dir=args.output_images_dir,
                         output_video=args.output_video,