update browse_dataset (#1817)

tools/misc/browse_dataset.py

@@ -2,28 +2,27 @@
 import argparse
 from os import path as osp
 
-import mmengine
-from mmengine import Config, DictAction, mkdir_or_exist
+from mmengine.config import Config, DictAction
+from mmengine.utils import ProgressBar, mkdir_or_exist
 
-from mmdet3d.datasets import build_dataset
-from mmdet3d.registry import VISUALIZERS
-from mmdet3d.utils import register_all_modules
+from mmdet3d.registry import DATASETS, VISUALIZERS
+from mmdet3d.utils import register_all_modules, replace_ceph_backend
 
 
 def parse_args():
     parser = argparse.ArgumentParser(description='Browse a dataset')
     parser.add_argument('config', help='train config file path')
-    parser.add_argument(
-        '--skip-type',
-        type=str,
-        nargs='+',
-        default=['Normalize'],
-        help='skip some useless pipeline')
     parser.add_argument(
         '--output-dir',
         default=None,
         type=str,
         help='If there is no display interface, you can save it')
+    parser.add_argument('--not-show', default=False, action='store_true')
+    parser.add_argument(
+        '--show-interval',
+        type=float,
+        default=2,
+        help='the interval of show (s)')
     parser.add_argument(
         '--task',
         type=str,
@@ -34,10 +33,7 @@ def parse_args():
         action='store_true',
         help='Whether to visualize augmented datasets or original dataset.')
     parser.add_argument(
-        '--online',
-        action='store_true',
-        help='Whether to perform online visualization. Note that you often '
-        'need a monitor to do so.')
+        '--ceph', action='store_true', help='Use ceph as data storage backend')
     parser.add_argument(
         '--cfg-options',
         nargs='+',
@@ -52,20 +48,22 @@ def parse_args():
     return args
 
 
-def build_data_cfg(config_path, skip_type, aug, cfg_options):
+def build_data_cfg(config_path, aug, cfg_options):
     """Build data config for loading visualization data."""
 
     cfg = Config.fromfile(config_path)
     if cfg_options is not None:
         cfg.merge_from_dict(cfg_options)
-    # extract inner dataset of `RepeatDataset` as `cfg.data.train`
-    # so we don't need to worry about it later
-    if cfg.data.train['type'] == 'RepeatDataset':
-        cfg.data.train = cfg.data.train.dataset
+
+    # extract inner dataset of `RepeatDataset` as
+    # `cfg.train_dataloader.dataset` so we don't
+    # need to worry about it later
+    if cfg.train_dataloader.dataset['type'] == 'RepeatDataset':
+        cfg.train_dataloader.dataset = cfg.train_dataloader.dataset.dataset
     # use only first dataset for `ConcatDataset`
-    if cfg.data.train['type'] == 'ConcatDataset':
-        cfg.data.train = cfg.data.train.datasets[0]
-    train_data_cfg = cfg.data.train
+    if cfg.train_dataloader.dataset['type'] == 'ConcatDataset':
+        cfg.train_dataloader.dataset = cfg.train_dataloader.dataset.datasets[0]
+    train_data_cfg = cfg.train_dataloader.dataset
 
     if aug:
         show_pipeline = cfg.train_pipeline
@@ -74,16 +72,14 @@ def build_data_cfg(config_path, skip_type, aug, cfg_options):
         for i in range(len(cfg.train_pipeline)):
             if cfg.train_pipeline[i]['type'] == 'LoadAnnotations3D':
                 show_pipeline.insert(i, cfg.train_pipeline[i])
-            # Collect points as well as labels
-            if cfg.train_pipeline[i]['type'] == 'Pack3DInputs':
-                if show_pipeline[-1]['type'] == 'Pack3DInputs':
+            # Collect data as well as labels
+            if cfg.train_pipeline[i]['type'] == 'Pack3DDetInputs':
+                if show_pipeline[-1]['type'] == 'Pack3DDetInputs':
                     show_pipeline[-1] = cfg.train_pipeline[i]
                 else:
                     show_pipeline.append(cfg.train_pipeline[i])
 
-    train_data_cfg['pipeline'] = [
-        x for x in show_pipeline if x['type'] not in skip_type
-    ]
+    train_data_cfg['pipeline'] = show_pipeline
 
     return cfg
 
@@ -94,18 +90,21 @@ def main():
     if args.output_dir is not None:
         mkdir_or_exist(args.output_dir)
 
-    cfg = build_data_cfg(args.config, args.skip_type, args.aug,
-                         args.cfg_options)
+    cfg = build_data_cfg(args.config, args.aug, args.cfg_options)
+
+    # TODO: We will unify the ceph support approach with other OpenMMLab repos
+    if args.ceph:
+        cfg = replace_ceph_backend(cfg)
 
     # register all modules in mmdet3d into the registries
     register_all_modules()
 
     try:
-        dataset = build_dataset(
+        dataset = DATASETS.build(
             cfg.train_dataloader.dataset,
             default_args=dict(filter_empty_gt=False))
     except TypeError:  # seg dataset doesn't have `filter_empty_gt` key
-        dataset = build_dataset(cfg.train_dataloader.dataset)
+        dataset = DATASETS.build(cfg.train_dataloader.dataset)
 
     # configure visualization mode
     vis_task = args.task  # 'det', 'seg', 'multi_modality-det', 'mono-det'
@@ -113,7 +112,7 @@ def main():
     visualizer = VISUALIZERS.build(cfg.visualizer)
     visualizer.dataset_meta = dataset.metainfo
 
-    progress_bar = mmengine.ProgressBar(len(dataset))
+    progress_bar = ProgressBar(len(dataset))
 
     for item in dataset:
         # the 3D Boxes in input could be in any of three coordinates
@@ -126,7 +125,7 @@ def main():
         visualizer.add_datasample(
             '3d visualzier',
             data_input,
-            data_sample,
+            data_sample=data_sample,
             show=not args.not_show,
             wait_time=args.show_interval,
             out_file=out_file,