mmpose/apis/inferencers/base_mmpose_inferencer.py

# Copyright (c) OpenMMLab. All rights reserved.
import inspect
import logging
import mimetypes
import os
from collections import defaultdict
from typing import (Callable, Dict, Generator, Iterable, List, Optional,
                    Sequence, Tuple, Union)

import cv2
import mmcv
import mmengine
import numpy as np
import torch.nn as nn
from mmengine.config import Config, ConfigDict
from mmengine.dataset import Compose
from mmengine.fileio import (get_file_backend, isdir, join_path,
                             list_dir_or_file)
from mmengine.infer.infer import BaseInferencer, ModelType
from mmengine.logging import print_log
from mmengine.registry import init_default_scope
from mmengine.runner.checkpoint import _load_checkpoint_to_model
from mmengine.structures import InstanceData
from mmengine.utils import mkdir_or_exist
from rich.progress import track

from mmpose.apis.inference import dataset_meta_from_config
from mmpose.registry import DATASETS
from mmpose.structures import PoseDataSample, split_instances
from .utils import default_det_models

try:
    from mmdet.apis.det_inferencer import DetInferencer
    has_mmdet = True
except (ImportError, ModuleNotFoundError):
    has_mmdet = False

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


class BaseMMPoseInferencer(BaseInferencer):
    """The base class for MMPose inferencers."""

    preprocess_kwargs: set = {'bbox_thr', 'nms_thr', 'bboxes'}
    forward_kwargs: set = set()
    visualize_kwargs: set = {
        'return_vis', 'show', 'wait_time', 'draw_bbox', 'radius', 'thickness',
        'kpt_thr', 'vis_out_dir', 'black_background'
    }
    postprocess_kwargs: set = {'pred_out_dir', 'return_datasample'}

    def __init__(self,
                 model: Union[ModelType, str, None] = None,
                 weights: Optional[str] = None,
                 device: Optional[str] = None,
                 scope: Optional[str] = None,
                 show_progress: bool = False) -> None:
        super().__init__(
            model, weights, device, scope, show_progress=show_progress)

    def _init_detector(
        self,
        det_model: Optional[Union[ModelType, str]] = None,
        det_weights: Optional[str] = None,
        det_cat_ids: Optional[Union[int, Tuple]] = None,
        device: Optional[str] = None,
    ):
        object_type = DATASETS.get(self.cfg.dataset_type).__module__.split(
            'datasets.')[-1].split('.')[0].lower()

        if det_model in ('whole_image', 'whole-image') or \
            (det_model is None and
                object_type not in default_det_models):
            self.detector = None

        else:
            det_scope = 'mmdet'
            if det_model is None:
                det_info = default_det_models[object_type]
                det_model, det_weights, det_cat_ids = det_info[
                    'model'], det_info['weights'], det_info['cat_ids']
            elif os.path.exists(det_model):
                det_cfg = Config.fromfile(det_model)
                det_scope = det_cfg.default_scope

            if has_mmdet:
                det_kwargs = dict(
                    model=det_model,
                    weights=det_weights,
                    device=device,
                    scope=det_scope,
                )
                # for compatibility with low version of mmdet
                if 'show_progress' in inspect.signature(
                        DetInferencer).parameters:
                    det_kwargs['show_progress'] = False

                self.detector = DetInferencer(**det_kwargs)
            else:
                raise RuntimeError(
                    'MMDetection (v3.0.0 or above) is required to build '
                    'inferencers for top-down pose estimation models.')

            if isinstance(det_cat_ids, (tuple, list)):
                self.det_cat_ids = det_cat_ids
            else:
                self.det_cat_ids = (det_cat_ids, )

    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)
            checkpoint_meta = checkpoint.get('meta', {})
            # save the dataset_meta in the model for convenience
            if 'dataset_meta' in checkpoint_meta:
                # mmpose 1.x
                model.dataset_meta = checkpoint_meta['dataset_meta']
            else:
                print_log(
                    'dataset_meta are not saved in the checkpoint\'s '
                    'meta data, load via config.',
                    logger='current',
                    level=logging.WARNING)
                model.dataset_meta = dataset_meta_from_config(
                    cfg, dataset_mode='train')
        else:
            print_log(
                'Checkpoint is not loaded, and the inference '
                'result is calculated by the randomly initialized '
                'model!',
                logger='current',
                level=logging.WARNING)
            model.dataset_meta = dataset_meta_from_config(
                cfg, dataset_mode='train')

    def _inputs_to_list(self, inputs: InputsType) -> Iterable:
        """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`.
        """
        self._video_input = False

        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
                filepath_list = [
                    join_path(inputs, fname)
                    for fname in list_dir_or_file(inputs, list_dir=False)
                ]
                inputs = []
                for filepath in filepath_list:
                    input_type = mimetypes.guess_type(filepath)[0].split(
                        '/')[0]
                    if input_type == 'image':
                        inputs.append(filepath)
                inputs.sort()
            else:
                # if inputs is a path to a video file, it will be converted
                # to a list containing separated frame filenames
                input_type = mimetypes.guess_type(inputs)[0].split('/')[0]
                if input_type == 'video':
                    self._video_input = True
                    video = mmcv.VideoReader(inputs)
                    self.video_info = dict(
                        fps=video.fps,
                        name=os.path.basename(inputs),
                        writer=None,
                        width=video.width,
                        height=video.height,
                        predictions=[])
                    inputs = video
                elif input_type == 'image':
                    inputs = [inputs]
                else:
                    raise ValueError(f'Expected input to be an image, video, '
                                     f'or folder, but received {inputs} of '
                                     f'type {input_type}.')

        elif isinstance(inputs, np.ndarray):
            inputs = [inputs]

        return inputs

    def _get_webcam_inputs(self, inputs: str) -> Generator:
        """Sets up and returns a generator function that reads frames from a
        webcam input. The generator function returns a new frame each time it
        is iterated over.

        Args:
            inputs (str): A string describing the webcam input, in the format
                "webcam:id".

        Returns:
            A generator function that yields frames from the webcam input.

        Raises:
            ValueError: If the inputs string is not in the expected format.
        """

        # Ensure the inputs string is in the expected format.
        inputs = inputs.lower()
        assert inputs.startswith('webcam'), f'Expected input to start with ' \
            f'"webcam", but got "{inputs}"'

        # Parse the camera ID from the inputs string.
        inputs_ = inputs.split(':')
        if len(inputs_) == 1:
            camera_id = 0
        elif len(inputs_) == 2 and str.isdigit(inputs_[1]):
            camera_id = int(inputs_[1])
        else:
            raise ValueError(
                f'Expected webcam input to have format "webcam:id", '
                f'but got "{inputs}"')

        # Attempt to open the video capture object.
        vcap = cv2.VideoCapture(camera_id)
        if not vcap.isOpened():
            print_log(
                f'Cannot open camera (ID={camera_id})',
                logger='current',
                level=logging.WARNING)
            return []

        # Set video input flag and metadata.
        self._video_input = True
        (major_ver, minor_ver, subminor_ver) = (cv2.__version__).split('.')
        if int(major_ver) < 3:
            fps = vcap.get(cv2.cv.CV_CAP_PROP_FPS)
            width = vcap.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)
            height = vcap.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)
        else:
            fps = vcap.get(cv2.CAP_PROP_FPS)
            width = vcap.get(cv2.CAP_PROP_FRAME_WIDTH)
            height = vcap.get(cv2.CAP_PROP_FRAME_HEIGHT)
        self.video_info = dict(
            fps=fps,
            name='webcam.mp4',
            writer=None,
            width=width,
            height=height,
            predictions=[])

        def _webcam_reader() -> Generator:
            while True:
                if cv2.waitKey(5) & 0xFF == 27:
                    vcap.release()
                    break

                ret_val, frame = vcap.read()
                if not ret_val:
                    break

                yield frame

        return _webcam_reader()

    def _init_pipeline(self, cfg: ConfigType) -> Callable:
        """Initialize the test pipeline.

        Args:
            cfg (ConfigType): model config path or dict

        Returns:
            A pipeline to handle various input data, such as ``str``,
            ``np.ndarray``. The returned pipeline will be used to process
            a single data.
        """
        scope = cfg.get('default_scope', 'mmpose')
        if scope is not None:
            init_default_scope(scope)
        return Compose(cfg.test_dataloader.dataset.pipeline)

    def update_model_visualizer_settings(self, **kwargs):
        """Update the settings of models and visualizer according to inference
        arguments."""

        pass

    def preprocess(self,
                   inputs: InputsType,
                   batch_size: int = 1,
                   bboxes: Optional[List] = None,
                   bbox_thr: float = 0.3,
                   nms_thr: float = 0.3,
                   **kwargs):
        """Process the inputs into a model-feedable format.

        Args:
            inputs (InputsType): Inputs given by user.
            batch_size (int): batch size. Defaults to 1.
            bbox_thr (float): threshold for bounding box detection.
                Defaults to 0.3.
            nms_thr (float): IoU threshold for bounding box NMS.
                Defaults to 0.3.

        Yields:
            Any: Data processed by the ``pipeline`` and ``collate_fn``.
            List[str or np.ndarray]: List of original inputs in the batch
        """

        # One-stage pose estimators perform prediction filtering within the
        # head's `predict` method. Here, we set the arguments for filtering
        if self.cfg.model.type == 'BottomupPoseEstimator':
            # 1. init with default arguments
            test_cfg = self.model.head.test_cfg.copy()
            # 2. update the score_thr and nms_thr in the test_cfg of the head
            if 'score_thr' in test_cfg:
                test_cfg['score_thr'] = bbox_thr
            if 'nms_thr' in test_cfg:
                test_cfg['nms_thr'] = nms_thr
            self.model.test_cfg = test_cfg

        for i, input in enumerate(inputs):
            bbox = bboxes[i] if bboxes else []
            data_infos = self.preprocess_single(
                input,
                index=i,
                bboxes=bbox,
                bbox_thr=bbox_thr,
                nms_thr=nms_thr,
                **kwargs)
            # only supports inference with batch size 1
            yield self.collate_fn(data_infos), [input]

    def __call__(
        self,
        inputs: InputsType,
        return_datasamples: bool = False,
        batch_size: int = 1,
        out_dir: Optional[str] = None,
        **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.
            out_dir (str, optional): directory to save visualization
                results and predictions. Will be overoden if vis_out_dir or
                pred_out_dir are given. Defaults to None
            **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.
        """
        if out_dir is not None:
            if 'vis_out_dir' not in kwargs:
                kwargs['vis_out_dir'] = f'{out_dir}/visualizations'
            if 'pred_out_dir' not in kwargs:
                kwargs['pred_out_dir'] = f'{out_dir}/predictions'

        (
            preprocess_kwargs,
            forward_kwargs,
            visualize_kwargs,
            postprocess_kwargs,
        ) = self._dispatch_kwargs(**kwargs)

        self.update_model_visualizer_settings(**kwargs)

        # preprocessing
        if isinstance(inputs, str) and inputs.startswith('webcam'):
            inputs = self._get_webcam_inputs(inputs)
            batch_size = 1
            if not visualize_kwargs.get('show', False):
                print_log(
                    'The display mode is closed when using webcam '
                    'input. It will be turned on automatically.',
                    logger='current',
                    level=logging.WARNING)
            visualize_kwargs['show'] = True
        else:
            inputs = self._inputs_to_list(inputs)

        # check the compatibility between inputs/outputs
        if not self._video_input and len(inputs) > 0:
            vis_out_dir = visualize_kwargs.get('vis_out_dir', None)
            if vis_out_dir is not None:
                _, file_extension = os.path.splitext(vis_out_dir)
                assert not file_extension, f'the argument `vis_out_dir` ' \
                    f'should be a folder while the input contains multiple ' \
                    f'images, but got {vis_out_dir}'

        if 'bbox_thr' in self.forward_kwargs:
            forward_kwargs['bbox_thr'] = preprocess_kwargs.get('bbox_thr', -1)
        inputs = self.preprocess(
            inputs, batch_size=batch_size, **preprocess_kwargs)

        preds = []

        for proc_inputs, ori_inputs in (track(inputs, description='Inference')
                                        if self.show_progress else inputs):
            preds = self.forward(proc_inputs, **forward_kwargs)

            visualization = self.visualize(ori_inputs, preds,
                                           **visualize_kwargs)
            results = self.postprocess(
                preds,
                visualization,
                return_datasamples=return_datasamples,
                **postprocess_kwargs)
            yield results

        if self._video_input:
            self._finalize_video_processing(
                postprocess_kwargs.get('pred_out_dir', ''))

        # In 3D Inferencers, some intermediate results (e.g. 2d keypoints)
        # will be temporarily stored in `self._buffer`. It's essential to
        # clear this information to prevent any interference with subsequent
        # inferences.
        if hasattr(self, '_buffer'):
            self._buffer.clear()

    def visualize(self,
                  inputs: list,
                  preds: List[PoseDataSample],
                  return_vis: bool = False,
                  show: bool = False,
                  draw_bbox: bool = False,
                  wait_time: float = 0,
                  radius: int = 3,
                  thickness: int = 1,
                  kpt_thr: float = 0.3,
                  vis_out_dir: str = '',
                  window_name: str = '',
                  black_background: bool = False,
                  **kwargs) -> List[np.ndarray]:
        """Visualize predictions.

        Args:
            inputs (list): Inputs preprocessed by :meth:`_inputs_to_list`.
            preds (Any): Predictions of the model.
            return_vis (bool): Whether to return images with predicted results.
            show (bool): Whether to display the image in a popup window.
                Defaults to False.
            wait_time (float): The interval of show (ms). Defaults to 0
            draw_bbox (bool): Whether to draw the bounding boxes.
                Defaults to False
            radius (int): Keypoint radius for visualization. Defaults to 3
            thickness (int): Link thickness for visualization. Defaults to 1
            kpt_thr (float): The threshold to visualize the keypoints.
                Defaults to 0.3
            vis_out_dir (str, optional): Directory to save visualization
                results w/o predictions. If left as empty, no file will
                be saved. Defaults to ''.
            window_name (str, optional): Title of display window.
            black_background (bool, optional): Whether to plot keypoints on a
                black image instead of the input image. Defaults to False.

        Returns:
            List[np.ndarray]: Visualization results.
        """
        if (not return_vis) and (not show) and (not vis_out_dir):
            return

        if getattr(self, 'visualizer', None) is None:
            raise ValueError('Visualization needs the "visualizer" term'
                             'defined in the config, but got None.')

        self.visualizer.radius = radius
        self.visualizer.line_width = thickness

        results = []

        for single_input, pred in zip(inputs, preds):
            if isinstance(single_input, str):
                img = mmcv.imread(single_input, channel_order='rgb')
            elif isinstance(single_input, np.ndarray):
                img = mmcv.bgr2rgb(single_input)
            else:
                raise ValueError('Unsupported input type: '
                                 f'{type(single_input)}')
            if black_background:
                img = img * 0

            img_name = os.path.basename(pred.metainfo['img_path'])
            window_name = window_name if window_name else img_name

            # since visualization and inference utilize the same process,
            # the wait time is reduced when a video input is utilized,
            # thereby eliminating the issue of inference getting stuck.
            wait_time = 1e-5 if self._video_input else wait_time

            visualization = self.visualizer.add_datasample(
                window_name,
                img,
                pred,
                draw_gt=False,
                draw_bbox=draw_bbox,
                show=show,
                wait_time=wait_time,
                kpt_thr=kpt_thr,
                **kwargs)
            results.append(visualization)

            if vis_out_dir:
                self.save_visualization(
                    visualization,
                    vis_out_dir,
                    img_name=img_name,
                )

        if return_vis:
            return results
        else:
            return []

    def save_visualization(self, visualization, vis_out_dir, img_name=None):
        out_img = mmcv.rgb2bgr(visualization)
        _, file_extension = os.path.splitext(vis_out_dir)
        if file_extension:
            dir_name = os.path.dirname(vis_out_dir)
            file_name = os.path.basename(vis_out_dir)
        else:
            dir_name = vis_out_dir
            file_name = None
        mkdir_or_exist(dir_name)

        if self._video_input:

            if self.video_info['writer'] is None:
                fourcc = cv2.VideoWriter_fourcc(*'mp4v')
                if file_name is None:
                    file_name = os.path.basename(self.video_info['name'])
                out_file = join_path(dir_name, file_name)
                self.video_info['output_file'] = out_file
                self.video_info['writer'] = cv2.VideoWriter(
                    out_file, fourcc, self.video_info['fps'],
                    (visualization.shape[1], visualization.shape[0]))
            self.video_info['writer'].write(out_img)

        else:
            if file_name is None:
                file_name = img_name if img_name else 'visualization.jpg'

            out_file = join_path(dir_name, file_name)
            mmcv.imwrite(out_img, out_file)
            print_log(
                f'the output image has been saved at {out_file}',
                logger='current',
                level=logging.INFO)

    def postprocess(
        self,
        preds: List[PoseDataSample],
        visualization: List[np.ndarray],
        return_datasample=None,
        return_datasamples=False,
        pred_out_dir: str = '',
    ) -> 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.

        Args:
            preds (List[Dict]): Predictions of the model.
            visualization (np.ndarray): Visualized predictions.
            return_datasamples (bool): Whether to return results as
                datasamples. Defaults to False.
            pred_out_dir (str): Directory to save the inference results w/o
                visualization. If left as empty, no file will be saved.
                Defaults to ''.

        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_datasamples=False``, it usually should be a
              json-serializable dict containing only basic data elements such
              as strings and numbers.
        """
        if return_datasample is not None:
            print_log(
                'The `return_datasample` argument is deprecated '
                'and will be removed in future versions. Please '
                'use `return_datasamples`.',
                logger='current',
                level=logging.WARNING)
            return_datasamples = return_datasample

        result_dict = defaultdict(list)

        result_dict['visualization'] = visualization
        for pred in preds:
            if not return_datasamples:
                # convert datasamples to list of instance predictions
                pred = split_instances(pred.pred_instances)
            result_dict['predictions'].append(pred)

        if pred_out_dir != '':
            for pred, data_sample in zip(result_dict['predictions'], preds):
                if self._video_input:
                    # For video or webcam input, predictions for each frame
                    # are gathered in the 'predictions' key of 'video_info'
                    # dictionary. All frame predictions are then stored into
                    # a single file after processing all frames.
                    self.video_info['predictions'].append(pred)
                else:
                    # For non-video inputs, predictions are stored in separate
                    # JSON files. The filename is determined by the basename
                    # of the input image path with a '.json' extension. The
                    # predictions are then dumped into this file.
                    fname = os.path.splitext(
                        os.path.basename(
                            data_sample.metainfo['img_path']))[0] + '.json'
                    mmengine.dump(
                        pred, join_path(pred_out_dir, fname), indent='  ')

        return result_dict

    def _finalize_video_processing(
        self,
        pred_out_dir: str = '',
    ):
        """Finalize video processing by releasing the video writer and saving
        predictions to a file.

        This method should be called after completing the video processing. It
        releases the video writer, if it exists, and saves the predictions to a
        JSON file if a prediction output directory is provided.
        """

        # Release the video writer if it exists
        if self.video_info['writer'] is not None:
            out_file = self.video_info['output_file']
            print_log(
                f'the output video has been saved at {out_file}',
                logger='current',
                level=logging.INFO)
            self.video_info['writer'].release()

        # Save predictions
        if pred_out_dir:
            fname = os.path.splitext(
                os.path.basename(self.video_info['name']))[0] + '.json'
            predictions = [
                dict(frame_id=i, instances=pred)
                for i, pred in enumerate(self.video_info['predictions'])
            ]

            mmengine.dump(
                predictions, join_path(pred_out_dir, fname), indent='  ')