submodule/sapiens/engine/mmengine/hub/hub.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import importlib
import os.path as osp

from mmengine.config import Config
from mmengine.config.utils import (_get_cfg_metainfo,
                                   _get_external_cfg_base_path,
                                   _get_package_and_cfg_path)
from mmengine.registry import MODELS, DefaultScope
from mmengine.runner import load_checkpoint
from mmengine.utils import get_installed_path, install_package


def get_config(cfg_path: str, pretrained: bool = False) -> Config:
    """Get config from external package.

    Args:
        cfg_path (str): External relative config path.
        pretrained (bool): Whether to save pretrained model path. If
            ``pretrained==True``, the url of pretrained model can be accessed
            by ``cfg.model_path``. Defaults to False.

    Examples:
        >>> cfg = get_config('mmdet::faster_rcnn/faster-rcnn_r50_fpn_1x_coco.py', pretrained=True)
        >>> # Equivalent to
        >>> # cfg = Config.fromfile('/path/to/faster-rcnn_r50_fpn_1x_coco.py')
        >>> cfg.model_path
        https://download.openmmlab.com/mmdetection/v2.0/faster_rcnn/faster_rcnn_r50_fpn_1x_coco/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth

    Returns:
        Config: A `Config` parsed from external package.
    """  # noqa E301
    # Get package name and relative config path.
    package, cfg_path = _get_package_and_cfg_path(cfg_path)
    # Install package if it's not installed.
    install_package(package)
    package_path = get_installed_path(package)
    try:
        # Use `cfg_path` to search target config file.
        cfg_meta = _get_cfg_metainfo(package_path, cfg_path)
        cfg_path = osp.join(package_path, '.mim', cfg_meta['Config'])
        cfg = Config.fromfile(cfg_path)
        if pretrained:
            assert 'Weights' in cfg_meta, ('Cannot find `Weights` in cfg_file'
                                           '.metafile.yml, please check the'
                                           'metafile')
            cfg.model_path = cfg_meta['Weights']
    except ValueError:
        # Since the base config does not contain a metafile, the absolute
        # config is `osp.join(package_path, cfg_path_prefix, cfg_name)`
        cfg_path = _get_external_cfg_base_path(package_path, cfg_path)
        cfg = Config.fromfile(cfg_path)
    except Exception as e:
        raise e
    return cfg


def get_model(cfg_path: str, pretrained: bool = False, **kwargs):
    """Get built model from external package.

    Args:
        cfg_path (str): External relative config path with prefix
            'package::' and without suffix.
        pretrained (bool): Whether to load pretrained model. Defaults to False.
        kwargs (dict): Default arguments to build model.

    Examples:
        >>> model = get_model('mmdet::faster_rcnn/faster-rcnn_r50_fpn_1x_coco.py', pretrained=True)
        >>> type(model)
        <class 'mmdet.models.detectors.faster_rcnn.FasterRCNN'>

    Returns:
        nn.Module: Built model.
    """  # noqa E301
    package = cfg_path.split('::')[0]
    with DefaultScope.overwrite_default_scope(package):  # type: ignore
        cfg = get_config(cfg_path, pretrained)
        if 'data_preprocessor' in cfg:
            cfg.model.data_preprocessor = cfg.data_preprocessor
        models_module = importlib.import_module(f'{package}.utils')
        models_module.register_all_modules()  # type: ignore
        model = MODELS.build(cfg.model, default_args=kwargs)
        if pretrained:
            load_checkpoint(model, cfg.model_path)
            # Hack to use pretrained weights.
            # If we do not set _is_init here, Runner will call
            # `model.init_weights()` to overwrite the pretrained model.
            model._is_init = True
        return model