[Enhance] Support modifying palette through configuration (#9445)

* fix order of palette

* fix order

* fix UT

* update

* fix comment

* fix UT

* update

demo/image_demo.py

@@ -20,8 +20,8 @@ def parse_args():
         '--device', default='cuda:0', help='Device used for inference')
     parser.add_argument(
         '--palette',
-        default='coco',
-        choices=['coco', 'voc', 'citys', 'random'],
+        default='none',
+        choices=['coco', 'voc', 'citys', 'random', 'none'],
         help='Color palette used for visualization')
     parser.add_argument(
         '--score-thr', type=float, default=0.3, help='bbox score threshold')

mmdet/apis/inference.py

@@ -1,4 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import copy
 import warnings
 from pathlib import Path
 from typing import Optional, Sequence, Union
@@ -11,6 +12,7 @@
 from mmengine.config import Config
 from mmengine.runner import load_checkpoint
 
+from mmdet.registry import DATASETS
 from ..evaluation import get_classes
 from ..registry import MODELS
 from ..structures import DetDataSample, SampleList
@@ -20,7 +22,7 @@
 def init_detector(
     config: Union[str, Path, Config],
     checkpoint: Optional[str] = None,
-    palette: str = 'coco',
+    palette: str = 'none',
     device: str = 'cuda:0',
     cfg_options: Optional[dict] = None,
 ) -> nn.Module:
@@ -34,7 +36,7 @@ def init_detector(
         palette (str): Color palette used for visualization. If palette
             is stored in checkpoint, use checkpoint's palette first, otherwise
             use externally passed palette. Currently, supports 'coco', 'voc',
-            'citys' and 'random'. Defaults to coco.
+            'citys' and 'random'. Defaults to none.
         device (str): The device where the anchors will be put on.
             Defaults to cuda:0.
         cfg_options (dict, optional): Options to override some settings in
@@ -53,10 +55,15 @@ def init_detector(
     elif 'init_cfg' in config.model.backbone:
         config.model.backbone.init_cfg = None
     model = MODELS.build(config.model)
-    if checkpoint is not None:
+    if checkpoint is None:
+        warnings.simplefilter('once')
+        warnings.warn('checkpoint is None, use COCO classes by default.')
+        model.dataset_meta = {'classes': get_classes('coco')}
+    else:
         checkpoint = load_checkpoint(model, checkpoint, map_location='cpu')
         # Weights converted from elsewhere may not have meta fields.
         checkpoint_meta = checkpoint.get('meta', {})
+
         # save the dataset_meta in the model for convenience
         if 'dataset_meta' in checkpoint_meta:
             # mmdet 3.x, all keys should be lowercase
@@ -67,16 +74,31 @@ def init_detector(
         elif 'CLASSES' in checkpoint_meta:
             # < mmdet 3.x
             classes = checkpoint_meta['CLASSES']
-            model.dataset_meta = {'classes': classes, 'palette': palette}
+            model.dataset_meta = {'classes': classes}
         else:
             warnings.simplefilter('once')
             warnings.warn(
                 'dataset_meta or class names are not saved in the '
                 'checkpoint\'s meta data, use COCO classes by default.')
-            model.dataset_meta = {
-                'classes': get_classes('coco'),
-                'palette': palette
-            }
+            model.dataset_meta = {'classes': get_classes('coco')}
+
+    # Priority:  args.palette -> config -> checkpoint
+    if palette != 'none':
+        model.dataset_meta['palette'] = palette
+    else:
+        test_dataset_cfg = copy.deepcopy(config.test_dataloader.dataset)
+        # lazy init. We only need the metainfo.
+        test_dataset_cfg['lazy_init'] = True
+        metainfo = DATASETS.build(test_dataset_cfg).metainfo
+        cfg_palette = metainfo.get('palette', None)
+        if cfg_palette is not None:
+            model.dataset_meta['palette'] = cfg_palette
+        else:
+            if 'palette' not in model.dataset_meta:
+                warnings.warn(
+                    'palette does not exist, random is used by default. '
+                    'You can also set the palette to customize.')
+                model.dataset_meta['palette'] = 'random'
 
     model.cfg = config  # save the config in the model for convenience
     model.to(device)