mmpose/apis/inferencers/pose3d_inferencer.py

# Copyright (c) OpenMMLab. All rights reserved.
import os
from collections import defaultdict
from functools import partial
from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union

import mmcv
import numpy as np
import torch
from mmengine.config import Config, ConfigDict
from mmengine.infer.infer import ModelType
from mmengine.model import revert_sync_batchnorm
from mmengine.registry import init_default_scope
from mmengine.structures import InstanceData

from mmpose.apis import (_track_by_iou, _track_by_oks, collate_pose_sequence,
                         convert_keypoint_definition, extract_pose_sequence)
from mmpose.registry import INFERENCERS
from mmpose.structures import PoseDataSample, merge_data_samples
from .base_mmpose_inferencer import BaseMMPoseInferencer
from .pose2d_inferencer import Pose2DInferencer

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]]


@INFERENCERS.register_module(name='pose-estimation-3d')
@INFERENCERS.register_module()
class Pose3DInferencer(BaseMMPoseInferencer):
    """The inferencer for 3D pose estimation.

    Args:
        model (str, optional): Pretrained 2D pose estimation algorithm.
            It's the path to the config file or the model name defined in
            metafile. For example, it could be:

            - model alias, e.g. ``'body'``,
            - config name, e.g. ``'simcc_res50_8xb64-210e_coco-256x192'``,
            - config path

            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 "mmpose".
        det_model (str, optional): Config path or alias of detection model.
            Defaults to None.
        det_weights (str, optional): Path to the checkpoints of detection
            model. Defaults to None.
        det_cat_ids (int or list[int], optional): Category id for
            detection model. Defaults to None.
        output_heatmaps (bool, optional): Flag to visualize predicted
            heatmaps. If set to None, the default setting from the model
            config will be used. Default is None.
    """

    preprocess_kwargs: set = {
        'bbox_thr', 'nms_thr', 'bboxes', 'use_oks_tracking', 'tracking_thr',
        'disable_norm_pose_2d'
    }
    forward_kwargs: set = {'disable_rebase_keypoint'}
    visualize_kwargs: set = {
        'return_vis',
        'show',
        'wait_time',
        'draw_bbox',
        'radius',
        'thickness',
        'num_instances',
        'kpt_thr',
        'vis_out_dir',
    }
    postprocess_kwargs: set = {'pred_out_dir', 'return_datasample'}

    def __init__(self,
                 model: Union[ModelType, str],
                 weights: Optional[str] = None,
                 pose2d_model: Optional[Union[ModelType, str]] = None,
                 pose2d_weights: Optional[str] = None,
                 device: Optional[str] = None,
                 scope: Optional[str] = 'mmpose',
                 det_model: Optional[Union[ModelType, str]] = None,
                 det_weights: Optional[str] = None,
                 det_cat_ids: Optional[Union[int, Tuple]] = None,
                 show_progress: bool = False) -> None:

        init_default_scope(scope)
        super().__init__(
            model=model,
            weights=weights,
            device=device,
            scope=scope,
            show_progress=show_progress)
        self.model = revert_sync_batchnorm(self.model)

        # assign dataset metainfo to self.visualizer
        self.visualizer.set_dataset_meta(self.model.dataset_meta)

        # initialize 2d pose estimator
        self.pose2d_model = Pose2DInferencer(
            pose2d_model if pose2d_model else 'human', pose2d_weights, device,
            scope, det_model, det_weights, det_cat_ids)

        # helper functions
        self._keypoint_converter = partial(
            convert_keypoint_definition,
            pose_det_dataset=self.pose2d_model.model.
            dataset_meta['dataset_name'],
            pose_lift_dataset=self.model.dataset_meta['dataset_name'],
        )

        self._pose_seq_extractor = partial(
            extract_pose_sequence,
            causal=self.cfg.test_dataloader.dataset.get('causal', False),
            seq_len=self.cfg.test_dataloader.dataset.get('seq_len', 1),
            step=self.cfg.test_dataloader.dataset.get('seq_step', 1))

        self._video_input = False
        self._buffer = defaultdict(list)

    def preprocess_single(self,
                          input: InputType,
                          index: int,
                          bbox_thr: float = 0.3,
                          nms_thr: float = 0.3,
                          bboxes: Union[List[List], List[np.ndarray],
                                        np.ndarray] = [],
                          use_oks_tracking: bool = False,
                          tracking_thr: float = 0.3,
                          disable_norm_pose_2d: bool = False):
        """Process a single input into a model-feedable format.

        Args:
            input (InputType): The input provided by the user.
            index (int): The index of the input.
            bbox_thr (float, optional): The threshold for bounding box
                detection. Defaults to 0.3.
            nms_thr (float, optional): The Intersection over Union (IoU)
                threshold for bounding box Non-Maximum Suppression (NMS).
                Defaults to 0.3.
            bboxes (Union[List[List], List[np.ndarray], np.ndarray]):
                The bounding boxes to use. Defaults to [].
            use_oks_tracking (bool, optional): A flag that indicates
                whether OKS-based tracking should be used. Defaults to False.
            tracking_thr (float, optional): The threshold for tracking.
                Defaults to 0.3.
            disable_norm_pose_2d (bool, optional): A flag that indicates
                whether 2D pose normalization should be used.
                Defaults to False.

        Yields:
            Any: The data processed by the pipeline and collate_fn.

        This method first calculates 2D keypoints using the provided
        pose2d_model. The method also performs instance matching, which
        can use either OKS-based tracking or IOU-based tracking.
        """

        # calculate 2d keypoints
        results_pose2d = next(
            self.pose2d_model(
                input,
                bbox_thr=bbox_thr,
                nms_thr=nms_thr,
                bboxes=bboxes,
                merge_results=False,
                return_datasamples=True))['predictions']

        for ds in results_pose2d:
            ds.pred_instances.set_field(
                (ds.pred_instances.bboxes[..., 2:] -
                 ds.pred_instances.bboxes[..., :2]).prod(-1), 'areas')

        if not self._video_input:
            height, width = results_pose2d[0].metainfo['ori_shape']

            # Clear the buffer if inputs are individual images to prevent
            # carryover effects from previous images
            self._buffer.clear()

        else:
            height = self.video_info['height']
            width = self.video_info['width']
        img_path = results_pose2d[0].metainfo['img_path']

        # instance matching
        if use_oks_tracking:
            _track = partial(_track_by_oks)
        else:
            _track = _track_by_iou

        for result in results_pose2d:
            track_id, self._buffer['results_pose2d_last'], _ = _track(
                result, self._buffer['results_pose2d_last'], tracking_thr)
            if track_id == -1:
                pred_instances = result.pred_instances.cpu().numpy()
                keypoints = pred_instances.keypoints
                if np.count_nonzero(keypoints[:, :, 1]) >= 3:
                    next_id = self._buffer.get('next_id', 0)
                    result.set_field(next_id, 'track_id')
                    self._buffer['next_id'] = next_id + 1
                else:
                    # If the number of keypoints detected is small,
                    # delete that person instance.
                    result.pred_instances.keypoints[..., 1] = -10
                    result.pred_instances.bboxes *= 0
                    result.set_field(-1, 'track_id')
            else:
                result.set_field(track_id, 'track_id')
        self._buffer['pose2d_results'] = merge_data_samples(results_pose2d)

        # convert keypoints
        results_pose2d_converted = [ds.cpu().numpy() for ds in results_pose2d]
        for ds in results_pose2d_converted:
            ds.pred_instances.keypoints = self._keypoint_converter(
                ds.pred_instances.keypoints)
        self._buffer['pose_est_results_list'].append(results_pose2d_converted)

        # extract and pad input pose2d sequence
        pose_results_2d = self._pose_seq_extractor(
            self._buffer['pose_est_results_list'],
            frame_idx=index if self._video_input else 0)
        causal = self.cfg.test_dataloader.dataset.get('causal', False)
        target_idx = -1 if causal else len(pose_results_2d) // 2

        stats_info = self.model.dataset_meta.get('stats_info', {})
        bbox_center = stats_info.get('bbox_center', None)
        bbox_scale = stats_info.get('bbox_scale', None)

        pose_results_2d_copy = []
        for pose_res in pose_results_2d:
            pose_res_copy = []
            for data_sample in pose_res:

                data_sample_copy = PoseDataSample()
                data_sample_copy.gt_instances = \
                    data_sample.gt_instances.clone()
                data_sample_copy.pred_instances = \
                    data_sample.pred_instances.clone()
                data_sample_copy.track_id = data_sample.track_id

                kpts = data_sample.pred_instances.keypoints
                bboxes = data_sample.pred_instances.bboxes
                keypoints = []
                for k in range(len(kpts)):
                    kpt = kpts[k]
                    if not disable_norm_pose_2d:
                        bbox = bboxes[k]
                        center = np.array([[(bbox[0] + bbox[2]) / 2,
                                            (bbox[1] + bbox[3]) / 2]])
                        scale = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
                        keypoints.append((kpt[:, :2] - center) / scale *
                                         bbox_scale + bbox_center)
                    else:
                        keypoints.append(kpt[:, :2])
                data_sample_copy.pred_instances.set_field(
                    np.array(keypoints), 'keypoints')
                pose_res_copy.append(data_sample_copy)

            pose_results_2d_copy.append(pose_res_copy)
        pose_sequences_2d = collate_pose_sequence(pose_results_2d_copy, True,
                                                  target_idx)
        if not pose_sequences_2d:
            return []

        data_list = []
        for i, pose_seq in enumerate(pose_sequences_2d):
            data_info = dict()

            keypoints_2d = pose_seq.pred_instances.keypoints
            keypoints_2d = np.squeeze(
                keypoints_2d,
                axis=0) if keypoints_2d.ndim == 4 else keypoints_2d

            T, K, C = keypoints_2d.shape

            data_info['keypoints'] = keypoints_2d
            data_info['keypoints_visible'] = np.ones((
                T,
                K,
            ),
                                                     dtype=np.float32)
            data_info['lifting_target'] = np.zeros((1, K, 3), dtype=np.float32)
            data_info['factor'] = np.zeros((T, ), dtype=np.float32)
            data_info['lifting_target_visible'] = np.ones((1, K, 1),
                                                          dtype=np.float32)
            data_info['camera_param'] = dict(w=width, h=height)

            data_info.update(self.model.dataset_meta)
            data_info = self.pipeline(data_info)
            data_info['data_samples'].set_field(
                img_path, 'img_path', field_type='metainfo')
            data_list.append(data_info)

        return data_list

    @torch.no_grad()
    def forward(self,
                inputs: Union[dict, tuple],
                disable_rebase_keypoint: bool = False):
        """Perform forward pass through the model and process the results.

        Args:
            inputs (Union[dict, tuple]): The inputs for the model.
            disable_rebase_keypoint (bool, optional): Flag to disable rebasing
                the height of the keypoints. Defaults to False.

        Returns:
            list: A list of data samples, each containing the model's output
                results.
        """
        pose_lift_results = self.model.test_step(inputs)

        # Post-processing of pose estimation results
        pose_est_results_converted = self._buffer['pose_est_results_list'][-1]
        for idx, pose_lift_res in enumerate(pose_lift_results):
            # Update track_id from the pose estimation results
            pose_lift_res.track_id = pose_est_results_converted[idx].get(
                'track_id', 1e4)

            # align the shape of output keypoints coordinates and scores
            keypoints = pose_lift_res.pred_instances.keypoints
            keypoint_scores = pose_lift_res.pred_instances.keypoint_scores
            if keypoint_scores.ndim == 3:
                pose_lift_results[idx].pred_instances.keypoint_scores = \
                    np.squeeze(keypoint_scores, axis=1)
            if keypoints.ndim == 4:
                keypoints = np.squeeze(keypoints, axis=1)

            # Invert x and z values of the keypoints
            keypoints = keypoints[..., [0, 2, 1]]
            keypoints[..., 0] = -keypoints[..., 0]
            keypoints[..., 2] = -keypoints[..., 2]

            # If rebase_keypoint_height is True, adjust z-axis values
            if not disable_rebase_keypoint:
                keypoints[..., 2] -= np.min(
                    keypoints[..., 2], axis=-1, keepdims=True)

            pose_lift_results[idx].pred_instances.keypoints = keypoints

        pose_lift_results = sorted(
            pose_lift_results, key=lambda x: x.get('track_id', 1e4))

        data_samples = [merge_data_samples(pose_lift_results)]
        return data_samples

    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,
                  num_instances: int = 1,
                  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
        det_kpt_color = self.pose2d_model.visualizer.kpt_color
        det_dataset_skeleton = self.pose2d_model.visualizer.skeleton
        det_dataset_link_color = self.pose2d_model.visualizer.link_color
        self.visualizer.det_kpt_color = det_kpt_color
        self.visualizer.det_dataset_skeleton = det_dataset_skeleton
        self.visualizer.det_dataset_link_color = det_dataset_link_color

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

            # 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

            if num_instances < 0:
                num_instances = len(pred.pred_instances)

            visualization = self.visualizer.add_datasample(
                window_name,
                img,
                data_sample=pred,
                det_data_sample=self._buffer['pose2d_results'],
                draw_gt=False,
                draw_bbox=draw_bbox,
                show=show,
                wait_time=wait_time,
                dataset_2d=self.pose2d_model.model.
                dataset_meta['dataset_name'],
                dataset_3d=self.model.dataset_meta['dataset_name'],
                kpt_thr=kpt_thr,
                num_instances=num_instances)
            results.append(visualization)

            if vis_out_dir:
                img_name = os.path.basename(pred.metainfo['img_path']) \
                    if 'img_path' in pred.metainfo else None
                self.save_visualization(
                    visualization,
                    vis_out_dir,
                    img_name=img_name,
                )

        if return_vis:
            return results
        else:
            return []