Source code for mmengine.infer.infer

# Copyright (c) OpenMMLab. All rights reserved.
import copy
import importlib
import os.path as osp
import re
import warnings
from abc import ABCMeta, abstractmethod
from datetime import datetime
from typing import (Any, Callable, Dict, Iterable, List, Optional, Sequence,
                    Tuple, Union)

import numpy as np
import torch
import torch.nn as nn
from rich.progress import track

from mmengine.config import Config, ConfigDict
from mmengine.config.utils import MODULE2PACKAGE
from mmengine.dataset import pseudo_collate
from mmengine.device import get_device
from mmengine.fileio import (get_file_backend, isdir, join_path,
                             list_dir_or_file, load)
from mmengine.logging import print_log
from mmengine.registry import FUNCTIONS, MODELS, VISUALIZERS, DefaultScope
from mmengine.runner.checkpoint import (_load_checkpoint,
from mmengine.structures import InstanceData
from mmengine.visualization import Visualizer

InstanceList = List[InstanceData]
InputType = Union[str, np.ndarray, torch.Tensor]
InputsType = Union[InputType, Sequence[InputType]]
ImgType = Union[np.ndarray, Sequence[np.ndarray]]
ResType = Union[Dict, List[Dict]]
ConfigType = Union[Config, ConfigDict]
ModelType = Union[dict, ConfigType, str]

class InferencerMeta(ABCMeta):
    """Check the legality of the inferencer.

    All Inferencers should not define duplicated keys for
    ``preprocess_kwargs``, ``forward_kwargs``, ``visualize_kwargs`` and

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        assert isinstance(self.preprocess_kwargs, set)
        assert isinstance(self.forward_kwargs, set)
        assert isinstance(self.visualize_kwargs, set)
        assert isinstance(self.postprocess_kwargs, set)

        all_kwargs = (
            self.preprocess_kwargs | self.forward_kwargs
            | self.visualize_kwargs | self.postprocess_kwargs)

        assert len(all_kwargs) == (
            len(self.preprocess_kwargs) + len(self.forward_kwargs) +
            len(self.visualize_kwargs) + len(self.postprocess_kwargs)), (
                f'Class define error! {self.__name__} should not '
                'define duplicated keys for `preprocess_kwargs`, '
                '`forward_kwargs`, `visualize_kwargs` and '
                '`postprocess_kwargs` are not allowed.')

[docs]class BaseInferencer(metaclass=InferencerMeta): """Base inferencer for downstream tasks. The BaseInferencer provides the standard workflow for inference as follows: 1. Preprocess the input data by :meth:`preprocess`. 2. Forward the data to the model by :meth:`forward`. ``BaseInferencer`` assumes the model inherits from :class:`mmengine.models.BaseModel` and will call `model.test_step` in :meth:`forward` by default. 3. Visualize the results by :meth:`visualize`. 4. Postprocess and return the results by :meth:`postprocess`. When we call the subclasses inherited from BaseInferencer (not overriding ``__call__``), the workflow will be executed in order. All subclasses of BaseInferencer could define the following class attributes for customization: - ``preprocess_kwargs``: The keys of the kwargs that will be passed to :meth:`preprocess`. - ``forward_kwargs``: The keys of the kwargs that will be passed to :meth:`forward` - ``visualize_kwargs``: The keys of the kwargs that will be passed to :meth:`visualize` - ``postprocess_kwargs``: The keys of the kwargs that will be passed to :meth:`postprocess` All attributes mentioned above should be a ``set`` of keys (strings), and each key should not be duplicated. Actually, :meth:`__call__` will dispatch all the arguments to the corresponding methods according to the ``xxx_kwargs`` mentioned above, therefore, the key in sets should be unique to avoid ambiguous dispatching. Warning: If subclasses defined the class attributes mentioned above with duplicated keys, an ``AssertionError`` will be raised during import process. Subclasses inherited from ``BaseInferencer`` should implement :meth:`_init_pipeline`, :meth:`visualize` and :meth:`postprocess`: - _init_pipeline: Return a callable object to preprocess the input data. - visualize: Visualize the results returned by :meth:`forward`. - postprocess: Postprocess the results returned by :meth:`forward` and :meth:`visualize`. Args: model (str, optional): Path to the config file or the model name defined in metafile. Take the `mmdet metafile <>`_ as an example, the `model` could be `retinanet_r18_fpn_1x_coco` or its alias. If model is not specified, user must provide the `weights` saved by MMEngine which contains the config string. Defaults to None. weights (str, optional): Path to the checkpoint. If it is not specified and model is a model name of metafile, the weights will be loaded from metafile. Defaults to None. device (str, optional): Device to run inference. If None, the available device will be automatically used. Defaults to None. scope (str, optional): The scope of the model. Defaults to None. show_progress (bool): Control whether to display the progress bar during the inference process. Defaults to True. `New in version 0.7.4.` Note: Since ``Inferencer`` could be used to infer batch data, `collate_fn` should be defined. If `collate_fn` is not defined in config file, the `collate_fn` will be `pseudo_collate` by default. """ # noqa: E501 preprocess_kwargs: set = set() forward_kwargs: set = set() visualize_kwargs: set = set() postprocess_kwargs: set = set() def __init__(self, model: Union[ModelType, str, None] = None, weights: Optional[str] = None, device: Optional[str] = None, scope: Optional[str] = None, show_progress: bool = True) -> None: if scope is None: default_scope = DefaultScope.get_current_instance() if default_scope is not None: scope = default_scope.scope_name self.scope = scope # Load config to cfg cfg: ConfigType if isinstance(model, str): if osp.isfile(model): cfg = Config.fromfile(model) else: # Load config and weights from metafile. If `weights` is # assigned, the weights defined in metafile will be ignored. cfg, _weights = self._load_model_from_metafile(model) if weights is None: weights = _weights elif isinstance(model, (Config, ConfigDict)): cfg = copy.deepcopy(model) elif isinstance(model, dict): cfg = copy.deepcopy(ConfigDict(model)) elif model is None: if weights is None: raise ValueError( 'If model is None, the weights must be specified since ' 'the config needs to be loaded from the weights') cfg = ConfigDict() else: raise TypeError('model must be a filepath or any ConfigType' f'object, but got {type(model)}') if device is None: device = get_device() self.model = self._init_model(cfg, weights, device) # type: ignore self.pipeline = self._init_pipeline(cfg) self.collate_fn = self._init_collate(cfg) self.visualizer = self._init_visualizer(cfg) self.cfg = cfg self.show_progress = show_progress def __call__( self, inputs: InputsType, return_datasamples: bool = False, batch_size: int = 1, **kwargs, ) -> dict: """Call the inferencer. Args: inputs (InputsType): Inputs for the inferencer. return_datasamples (bool): Whether to return results as :obj:`BaseDataElement`. Defaults to False. batch_size (int): Batch size. Defaults to 1. **kwargs: Key words arguments passed to :meth:`preprocess`, :meth:`forward`, :meth:`visualize` and :meth:`postprocess`. Each key in kwargs should be in the corresponding set of ``preprocess_kwargs``, ``forward_kwargs``, ``visualize_kwargs`` and ``postprocess_kwargs``. Returns: dict: Inference and visualization results. """ ( preprocess_kwargs, forward_kwargs, visualize_kwargs, postprocess_kwargs, ) = self._dispatch_kwargs(**kwargs) ori_inputs = self._inputs_to_list(inputs) inputs = self.preprocess( ori_inputs, batch_size=batch_size, **preprocess_kwargs) preds = [] for data in (track(inputs, description='Inference') if self.show_progress else inputs): preds.extend(self.forward(data, **forward_kwargs)) visualization = self.visualize( ori_inputs, preds, **visualize_kwargs) # type: ignore # noqa: E501 results = self.postprocess(preds, visualization, return_datasamples, **postprocess_kwargs) return results def _inputs_to_list(self, inputs: InputsType) -> list: """Preprocess the inputs to a list. Preprocess inputs to a list according to its type: - list or tuple: return inputs - str: - Directory path: return all files in the directory - other cases: return a list containing the string. The string could be a path to file, a url or other types of string according to the task. Args: inputs (InputsType): Inputs for the inferencer. Returns: list: List of input for the :meth:`preprocess`. """ if isinstance(inputs, str): backend = get_file_backend(inputs) if hasattr(backend, 'isdir') and isdir(inputs): # Backends like HttpsBackend do not implement `isdir`, so only # those backends that implement `isdir` could accept the inputs # as a directory filename_list = list_dir_or_file(inputs, list_dir=False) inputs = [ join_path(inputs, filename) for filename in filename_list ] if not isinstance(inputs, (list, tuple)): inputs = [inputs] return list(inputs)
[docs] def preprocess(self, inputs: InputsType, batch_size: int = 1, **kwargs): """Process the inputs into a model-feedable format. Customize your preprocess by overriding this method. Preprocess should return an iterable object, of which each item will be used as the input of ``model.test_step``. ``BaseInferencer.preprocess`` will return an iterable chunked data, which will be used in __call__ like this: .. code-block:: python def __call__(self, inputs, batch_size=1, **kwargs): chunked_data = self.preprocess(inputs, batch_size, **kwargs) for batch in chunked_data: preds = self.forward(batch, **kwargs) Args: inputs (InputsType): Inputs given by user. batch_size (int): batch size. Defaults to 1. Yields: Any: Data processed by the ``pipeline`` and ``collate_fn``. """ chunked_data = self._get_chunk_data( map(self.pipeline, inputs), batch_size) yield from map(self.collate_fn, chunked_data)
[docs] @torch.no_grad() def forward(self, inputs: Union[dict, tuple], **kwargs) -> Any: """Feed the inputs to the model.""" return self.model.test_step(inputs)
[docs] @abstractmethod def visualize(self, inputs: list, preds: Any, show: bool = False, **kwargs) -> List[np.ndarray]: """Visualize predictions. Customize your visualization by overriding this method. visualize should return visualization results, which could be np.ndarray or any other objects. Args: inputs (list): Inputs preprocessed by :meth:`_inputs_to_list`. preds (Any): Predictions of the model. show (bool): Whether to display the image in a popup window. Defaults to False. Returns: List[np.ndarray]: Visualization results. """
[docs] @abstractmethod def postprocess( self, preds: Any, visualization: List[np.ndarray], return_datasample=False, **kwargs, ) -> dict: """Process the predictions and visualization results from ``forward`` and ``visualize``. This method should be responsible for the following tasks: 1. Convert datasamples into a json-serializable dict if needed. 2. Pack the predictions and visualization results and return them. 3. Dump or log the predictions. Customize your postprocess by overriding this method. Make sure ``postprocess`` will return a dict with visualization results and inference results. Args: preds (List[Dict]): Predictions of the model. visualization (np.ndarray): Visualized predictions. return_datasample (bool): Whether to return results as datasamples. Defaults to False. Returns: dict: Inference and visualization results with key ``predictions`` and ``visualization`` - ``visualization (Any)``: Returned by :meth:`visualize` - ``predictions`` (dict or DataSample): Returned by :meth:`forward` and processed in :meth:`postprocess`. If ``return_datasample=False``, it usually should be a json-serializable dict containing only basic data elements such as strings and numbers. """
def _load_model_from_metafile(self, model: str) -> Tuple[Config, str]: """Load config and weights from metafile. Args: model (str): model name defined in metafile. Returns: Tuple[Config, str]: Loaded Config and weights path defined in metafile. """ model = model.lower() assert self.scope is not None, ( 'scope should be initialized if you want ' 'to load config from metafile.') assert self.scope in MODULE2PACKAGE, ( f'{self.scope} not in {MODULE2PACKAGE}!,' 'please pass a valid scope.') repo_or_mim_dir = BaseInferencer._get_repo_or_mim_dir(self.scope) for model_cfg in BaseInferencer._get_models_from_metafile( repo_or_mim_dir): model_name = model_cfg['Name'].lower() model_aliases = model_cfg.get('Alias', []) if isinstance(model_aliases, str): model_aliases = [model_aliases.lower()] else: model_aliases = [alias.lower() for alias in model_aliases] if (model_name == model or model in model_aliases): cfg = Config.fromfile( osp.join(repo_or_mim_dir, model_cfg['Config'])) weights = model_cfg['Weights'] weights = weights[0] if isinstance(weights, list) else weights return cfg, weights raise ValueError(f'Cannot find model: {model} in {self.scope}') @staticmethod def _get_repo_or_mim_dir(scope): """Get the directory where the ``Configs`` located when the package is installed or ``PYTHONPATH`` is set. Args: scope (str): The scope of repository. Returns: str: The directory where the ``Configs`` is located. """ try: module = importlib.import_module(scope) except ImportError: if scope not in MODULE2PACKAGE: raise KeyError( f'{scope} is not a valid scope. The available scopes ' f'are {MODULE2PACKAGE.keys()}') else: project = MODULE2PACKAGE[scope] raise ImportError( f'Cannot import {scope} correctly, please try to install ' f'the {project} by "pip install {project}"') # Since none of OpenMMLab series packages are namespace packages # (, # The first element of module.__path__ means package installation path. package_path = module.__path__[0] if osp.exists(osp.join(osp.dirname(package_path), 'configs')): repo_dir = osp.dirname(package_path) return repo_dir else: mim_dir = osp.join(package_path, '.mim') if not osp.exists(osp.join(mim_dir, 'configs')): raise FileNotFoundError( f'Cannot find `configs` directory in {package_path}!, ' f'please check the completeness of the {scope}.') return mim_dir def _init_model( self, cfg: ConfigType, weights: Optional[str], device: str = 'cpu', ) -> nn.Module: """Initialize the model with the given config and checkpoint on the specific device. Args: cfg (ConfigType): Config containing the model information. weights (str, optional): Path to the checkpoint. device (str, optional): Device to run inference. Defaults to 'cpu'. Returns: nn.Module: Model loaded with checkpoint. """ checkpoint: Optional[dict] = None if weights is not None: checkpoint = _load_checkpoint(weights, map_location='cpu') if not cfg: assert checkpoint is not None try: # Prefer to get config from `message_hub` since `message_hub` # is a more stable module to store all runtime information. # However, the early version of MMEngine will not save config # in `message_hub`, so we will try to load config from `meta`. cfg_string = checkpoint['message_hub']['runtime_info']['cfg'] except KeyError: assert 'meta' in checkpoint, ( 'If model(config) is not provided, the checkpoint must' 'contain the config string in `meta` or `message_hub`, ' 'but both `meta` and `message_hub` are not found in the ' 'checkpoint.') meta = checkpoint['meta'] if 'cfg' in meta: cfg_string = meta['cfg'] else: raise ValueError( 'Cannot find the config in the checkpoint.') cfg.update( Config.fromstring(cfg_string, file_format='.py')._cfg_dict) # Delete the `pretrained` field to prevent model from loading the # the pretrained weights unnecessarily. if cfg.model.get('pretrained') is not None: del cfg.model.pretrained model = model.cfg = cfg self._load_weights_to_model(model, checkpoint, cfg) model.eval() return model def _load_weights_to_model(self, model: nn.Module, checkpoint: Optional[dict], cfg: Optional[ConfigType]) -> None: """Loading model weights and meta information from cfg and checkpoint. Subclasses could override this method to load extra meta information from ``checkpoint`` and ``cfg`` to model. Args: model (nn.Module): Model to load weights and meta information. checkpoint (dict, optional): The loaded checkpoint. cfg (Config or ConfigDict, optional): The loaded config. """ if checkpoint is not None: _load_checkpoint_to_model(model, checkpoint) else: warnings.warn('Checkpoint is not loaded, and the inference ' 'result is calculated by the randomly initialized ' 'model!') def _init_collate(self, cfg: ConfigType) -> Callable: """Initialize the ``collate_fn`` with the given config. The returned ``collate_fn`` will be used to collate the batch data. If will be used in :meth:`preprocess` like this .. code-block:: python def preprocess(self, inputs, batch_size, **kwargs): ... dataloader = map(self.collate_fn, dataloader) yield from dataloader Args: cfg (ConfigType): Config which could contained the `collate_fn` information. If `collate_fn` is not defined in config, it will be :func:`pseudo_collate`. Returns: Callable: Collate function. """ try: with FUNCTIONS.switch_scope_and_registry(self.scope) as registry: collate_fn = registry.get(cfg.test_dataloader.collate_fn) except AttributeError: collate_fn = pseudo_collate return collate_fn # type: ignore @abstractmethod def _init_pipeline(self, cfg: ConfigType) -> Callable: """Initialize the test pipeline. Return a pipeline to handle various input data, such as ``str``, ``np.ndarray``. It is an abstract method in BaseInferencer, and should be implemented in subclasses. The returned pipeline will be used to process a single data. It will be used in :meth:`preprocess` like this: .. code-block:: python def preprocess(self, inputs, batch_size, **kwargs): ... dataset = map(self.pipeline, dataset) ... """ def _init_visualizer(self, cfg: ConfigType) -> Optional[Visualizer]: """Initialize visualizers. Args: cfg (ConfigType): Config containing the visualizer information. Returns: Visualizer or None: Visualizer initialized with config. """ if 'visualizer' not in cfg: return None timestamp = str(datetime.timestamp( name = cfg.visualizer.get('name', timestamp) if Visualizer.check_instance_created(name): name = f'{name}-{timestamp}' = name return def _get_chunk_data(self, inputs: Iterable, chunk_size: int): """Get batch data from dataset. Args: inputs (Iterable): An iterable dataset. chunk_size (int): Equivalent to batch size. Yields: list: batch data. """ inputs_iter = iter(inputs) while True: try: chunk_data = [] for _ in range(chunk_size): processed_data = next(inputs_iter) chunk_data.append(processed_data) yield chunk_data except StopIteration: if chunk_data: yield chunk_data break def _dispatch_kwargs(self, **kwargs) -> Tuple[Dict, Dict, Dict, Dict]: """Dispatch kwargs to preprocess(), forward(), visualize() and postprocess() according to the actual demands. Returns: Tuple[Dict, Dict, Dict, Dict]: kwargs passed to preprocess, forward, visualize and postprocess respectively. """ # Ensure each argument only matches one function method_kwargs = self.preprocess_kwargs | self.forward_kwargs | \ self.visualize_kwargs | self.postprocess_kwargs union_kwargs = method_kwargs | set(kwargs.keys()) if union_kwargs != method_kwargs: unknown_kwargs = union_kwargs - method_kwargs raise ValueError( f'unknown argument {unknown_kwargs} for `preprocess`, ' '`forward`, `visualize` and `postprocess`') preprocess_kwargs = {} forward_kwargs = {} visualize_kwargs = {} postprocess_kwargs = {} for key, value in kwargs.items(): if key in self.preprocess_kwargs: preprocess_kwargs[key] = value elif key in self.forward_kwargs: forward_kwargs[key] = value elif key in self.visualize_kwargs: visualize_kwargs[key] = value else: postprocess_kwargs[key] = value return ( preprocess_kwargs, forward_kwargs, visualize_kwargs, postprocess_kwargs, ) @staticmethod def _get_models_from_metafile(dir: str): """Load model config defined in metafile from package path. Args: dir (str): Path to the directory of Config. It requires the directory ``Config``, file ``model-index.yml`` exists in the ``dir``. Yields: dict: Model config defined in metafile. """ meta_indexes = load(osp.join(dir, 'model-index.yml')) for meta_path in meta_indexes['Import']: # meta_path example: mmcls/.mim/configs/conformer/metafile.yml meta_path = osp.join(dir, meta_path) metainfo = load(meta_path) yield from metainfo['Models']
[docs] @staticmethod def list_models(scope: Optional[str] = None, patterns: str = r'.*'): """List models defined in metafile of corresponding packages. Args: scope (str, optional): The scope to which the model belongs. Defaults to None. patterns (str, optional): Regular expressions for the searched models. Once matched with ``Alias`` or ``Name`` filed in metafile, corresponding model will be added to the return list. Defaults to '.*'. Returns: dict: Model dict with model name and its alias. """ matched_models = [] if scope is None: default_scope = DefaultScope.get_current_instance() assert default_scope is not None, ( 'scope should be initialized if you want ' 'to load config from metafile.') assert scope in MODULE2PACKAGE, ( f'{scope} not in {MODULE2PACKAGE}!, please make pass a valid ' 'scope.') root_or_mim_dir = BaseInferencer._get_repo_or_mim_dir(scope) for model_cfg in BaseInferencer._get_models_from_metafile( root_or_mim_dir): model_name = [model_cfg['Name']] model_name.extend(model_cfg.get('Alias', [])) for name in model_name: if re.match(patterns, name) is not None: matched_models.append(name) output_str = '' for name in matched_models: output_str += f'model_name: {name}\n' print_log(output_str, logger='current') return matched_models

© Copyright 2022, mmengine contributors. Revision 66fb81f7.

Built with Sphinx using a theme provided by Read the Docs.
Read the Docs v: latest
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.