mmpose/apis/inferencers/base_mmpose_inferencer.py

# Copyright (c) OpenMMLab. All rights reserved.
import mimetypes
import os
import warnings
from collections import defaultdict
from typing import (Callable, Dict, Generator, Iterable, List, Optional,
                    Sequence, 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
from mmengine.runner.checkpoint import _load_checkpoint_to_model
from mmengine.structures import InstanceData
from mmengine.utils import mkdir_or_exist

from mmpose.apis.inference import dataset_meta_from_config
from mmpose.structures import PoseDataSample, split_instances

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',
    }
    postprocess_kwargs: set = {'pred_out_dir'}

    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:
                warnings.warn(
                    'dataset_meta are not saved in the checkpoint\'s '
                    'meta data, load via config.')
                model.dataset_meta = dataset_meta_from_config(
                    cfg, dataset_mode='train')
        else:
            warnings.warn('Checkpoint is not loaded, and the inference '
                          'result is calculated by the randomly initialized '
                          'model!')
            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,
                        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.
        """
        assert getattr(self.visualizer, 'backend', None) == 'opencv', \
            'Visualizer must utilize the OpenCV backend in order to ' \
            'support webcam inputs.'

        # 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():
            warnings.warn(f'Cannot open camera (ID={camera_id})')
            return []

        # Set video input flag and metadata.
        self._video_input = True
        self.video_info = dict(
            fps=10, name='webcam.mp4', writer=None, 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 _visualization_window_on_close(self, event):
        self._window_closing = True

    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.
        """
        return Compose(cfg.test_dataloader.dataset.pipeline)

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

        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``.
            List[str or np.ndarray]: List of original inputs in the batch
        """

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

    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 = '',
                  window_close_event_handler: Optional[Callable] = None
                  ) -> 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.
            window_close_event_handler (callable, optional):

        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)}')

            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,
                draw_heatmap=True,
                show=show,
                wait_time=wait_time,
                kpt_thr=kpt_thr)
            results.append(visualization)

            if vis_out_dir:
                out_img = mmcv.rgb2bgr(visualization)

                if self._video_input:

                    if self.video_info['writer'] is None:
                        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
                        mkdir_or_exist(vis_out_dir)
                        out_file = join_path(
                            vis_out_dir,
                            os.path.basename(self.video_info['name']))
                        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:
                    out_file = join_path(vis_out_dir, img_name)
                    mmcv.imwrite(out_img, out_file)

        if return_vis:
            return results
        else:
            return []

    def postprocess(
        self,
        preds: List[PoseDataSample],
        visualization: List[np.ndarray],
        return_datasample=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_datasample (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_datasample=False``, it usually should be a
              json-serializable dict containing only basic data elements such
              as strings and numbers.
        """

        result_dict = defaultdict(list)

        result_dict['visualization'] = visualization
        for pred in preds:
            if not return_datasample:
                # 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:
                    self.video_info['predictions'].append(pred)
                else:
                    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:
            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='  ')