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,
_load_checkpoint_to_model)
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
``postprocess_kwargs``.
"""
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 <https://github.com/open-mmlab/mmdetection/blob/master/configs/retinanet/metafile.yml>`_
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
# (https://docs.python.org/3/glossary.html#term-namespace-package),
# 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 = MODELS.build(cfg.model)
model.cfg = cfg
self._load_weights_to_model(model, checkpoint, cfg)
model.to(device)
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(datetime.now()))
name = cfg.visualizer.get('name', timestamp)
if Visualizer.check_instance_created(name):
name = f'{name}-{timestamp}'
cfg.visualizer.name = name
return VISUALIZERS.build(cfg.visualizer)
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