[Tools] Support replacing the ${key} with the value of cfg.key

mmdet/utils/__init__.py

@@ -2,10 +2,11 @@
 from .collect_env import collect_env
 from .logger import get_caller_name, get_root_logger, log_img_scale
 from .misc import find_latest_checkpoint, update_data_root
+from .replace import replace_config
 from .setup_env import setup_multi_processes
 
 __all__ = [
     'get_root_logger', 'collect_env', 'find_latest_checkpoint',
     'update_data_root', 'setup_multi_processes', 'get_caller_name',
-    'log_img_scale'
+    'log_img_scale', 'replace_config'
 ]

mmdet/utils/replace.py

@@ -0,0 +1,65 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import os.path as osp
+import re
+
+from mmcv.utils import Config
+
+
+def replace_config(ori_cfg):
+    """Replace the string "${key}" with the corresponding value.
+
+    Replace the "${key}" with the value of ori_cfg.key in the config. And
+    support replacing the chained ${key}. Such as, replace "${key0.key1}"
+    with the value of cfg.key0.key1. Code is modified from `vars.py
+    < https://github.com/microsoft/SoftTeacher/blob/main/ssod/utils/vars.py>`_  # noqa: E501
+
+    Args:
+        ori_cfg (mmcv.utils.config.Config):
+            The origin config with "${key}" generated from a file.
+
+    Returns:
+        replaced_cfg (mmcv.utils.config.Config):
+            The config with "${key}" replaced by the corresponding value.
+    """
+
+    def get_value(cfg, key):
+        chained_key = key.split('.')
+        if len(chained_key) == 1:
+            return cfg[chained_key[0]]
+        else:
+            return get_value(cfg[chained_key[0]], '.'.join(chained_key[1:]))
+
+    def replace_value(cfg):
+        if isinstance(cfg, dict):
+            return {key: replace_value(value) for key, value in cfg.items()}
+        elif isinstance(cfg, list):
+            return [replace_value(item) for item in cfg]
+        elif isinstance(cfg, tuple):
+            return tuple([replace_value(item) for item in cfg])
+        elif isinstance(cfg, str):
+            # replace the "${key}" with cfg.key
+            keys = pattern_key.findall(cfg)
+            values = [get_value(ori_cfg_dict, key[2:-1]) for key in keys]
+            # only support replacing one "${key}" for dict, list, or tuple
+            for key, value in zip(keys, values):
+                cfg = value if len(keys) == 1 else cfg.replace(key, str(value))
+            return cfg
+        else:
+            return cfg
+
+    # the pattern of string "${key}",
+    # which will be replaced by its value, such as "${model}"
+    pattern_key = re.compile(r'\$\{[a-zA-Z\d_.]*\}')
+    # ori_cfg is the cfg before being replaced
+    ori_cfg_dict = ori_cfg._cfg_dict.to_dict()
+    # cfg_name is part of the default working path
+    # work_dirs/${cfg_name}/${percent}/${fold}
+    ori_cfg_dict['cfg_name'] = osp.splitext(osp.basename(ori_cfg.filename))[0]
+    replaced_cfg = Config(
+        replace_value(ori_cfg_dict), filename=ori_cfg.filename)
+    # replace the model with semi_wrapper
+    if replaced_cfg.get('semi_wrapper', None) is not None:
+        replaced_cfg.model = replaced_cfg.semi_wrapper
+        replaced_cfg.pop('semi_wrapper')
+
+    return replaced_cfg

tests/test_utils/test_replace.py

@@ -0,0 +1,56 @@
+import os.path as osp
+import tempfile
+from copy import deepcopy
+
+from mmcv.utils import Config
+
+from mmdet.utils import replace_config
+
+
+def test_replace_config():
+    temp_file = tempfile.NamedTemporaryFile()
+    config_path = f'{temp_file.name}.py'
+    with open(config_path, 'w') as f:
+        f.write('configs')
+
+    ori_cfg_dict = dict()
+    ori_cfg_dict['work_dir'] = 'work_dirs/${cfg_name}/${percent}/${fold}'
+    ori_cfg_dict['percent'] = 5
+    ori_cfg_dict['fold'] = 1
+    ori_cfg_dict['semi_wrapper'] = dict(type='SoftTeacher', model='${model}')
+    ori_cfg_dict['model'] = dict(
+        type='FasterRCNN',
+        backbone=dict(type='ResNet'),
+        neck=dict(type='FPN'),
+        rpn_head=dict(type='RPNHead'),
+        roi_head=dict(type='StandardRoIHead'),
+        train_cfg=dict(
+            rpn=dict(
+                assigner=dict(type='MaxIoUAssigner'),
+                sampler=dict(type='RandomSampler'),
+            ),
+            rpn_proposal=dict(nms=dict(type='nms', iou_threshold=0.7)),
+            rcnn=dict(
+                assigner=dict(type='MaxIoUAssigner'),
+                sampler=dict(type='RandomSampler'),
+            ),
+        ),
+        test_cfg=dict(
+            rpn=dict(nms=dict(type='nms', iou_threshold=0.7)),
+            rcnn=dict(nms=dict(type='nms', iou_threshold=0.5)),
+        ),
+    )
+    ori_cfg_dict['iou_threshold'] = dict(
+        rpn_proposal_nms='${model.train_cfg.rpn_proposal.nms.iou_threshold}',
+        test_rpn_nms='${model.test_cfg.rpn.nms.iou_threshold}',
+        test_rcnn_nms='${model.test_cfg.rcnn.nms.iou_threshold}',
+    )
+
+    ori_cfg = Config(ori_cfg_dict, filename=config_path)
+    replaced_cfg = replace_config(deepcopy(ori_cfg))
+
+    assert replaced_cfg.work_dir \
+        == f'work_dirs/{osp.basename(temp_file.name)}/5/1'
+    assert replaced_cfg.model.model == ori_cfg.model
+    assert replaced_cfg.iou_threshold.rpn_proposal_nms \
+        == ori_cfg.model.train_cfg.rpn_proposal.nms.iou_threshold

tools/train.py

@@ -17,8 +17,8 @@
 from mmdet.apis import init_random_seed, set_random_seed, train_detector
 from mmdet.datasets import build_dataset
 from mmdet.models import build_detector
-from mmdet.utils import (collect_env, get_root_logger, setup_multi_processes,
-                         update_data_root)
+from mmdet.utils import (collect_env, get_root_logger, replace_config,
+                         setup_multi_processes, update_data_root)
 
 
 def parse_args():
@@ -142,6 +142,10 @@ def main():
         # use config filename as default work_dir if cfg.work_dir is None
         cfg.work_dir = osp.join('./work_dirs',
                                 osp.splitext(osp.basename(args.config))[0])
+
+    # replace the ${key} with the value of cfg.key
+    cfg = replace_config(cfg)
+
     if args.resume_from is not None:
         cfg.resume_from = args.resume_from
     cfg.auto_resume = args.auto_resume