Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

third_party/hamer/third-party/ViTPose/mmpose/models/detectors/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from .associative_embedding import AssociativeEmbedding
+from .interhand_3d import Interhand3D
+from .mesh import ParametricMesh
+from .multi_task import MultiTask
+from .multiview_pose import (DetectAndRegress, VoxelCenterDetector,
+                             VoxelSinglePose)
+from .pose_lifter import PoseLifter
+from .posewarper import PoseWarper
+from .top_down import TopDown
+from .top_down_moe import TopDownMoE
+
+__all__ = [
+    'TopDown', 'AssociativeEmbedding', 'ParametricMesh', 'MultiTask',
+    'PoseLifter', 'Interhand3D', 'PoseWarper', 'DetectAndRegress',
+    'VoxelCenterDetector', 'VoxelSinglePose', 'TopDownMoE'
+]

third_party/hamer/third-party/ViTPose/mmpose/models/detectors/associative_embedding.py

@@ -0,0 +1,420 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import warnings
+
+import mmcv
+import torch
+from mmcv.image import imwrite
+from mmcv.utils.misc import deprecated_api_warning
+from mmcv.visualization.image import imshow
+
+from mmpose.core.evaluation import (aggregate_scale, aggregate_stage_flip,
+                                    flip_feature_maps, get_group_preds,
+                                    split_ae_outputs)
+from mmpose.core.post_processing.group import HeatmapParser
+from mmpose.core.visualization import imshow_keypoints
+from .. import builder
+from ..builder import POSENETS
+from .base import BasePose
+
+try:
+    from mmcv.runner import auto_fp16
+except ImportError:
+    warnings.warn('auto_fp16 from mmpose will be deprecated from v0.15.0'
+                  'Please install mmcv>=1.1.4')
+    from mmpose.core import auto_fp16
+
+
+@POSENETS.register_module()
+class AssociativeEmbedding(BasePose):
+    """Associative embedding pose detectors.
+
+    Args:
+        backbone (dict): Backbone modules to extract feature.
+        keypoint_head (dict): Keypoint head to process feature.
+        train_cfg (dict): Config for training. Default: None.
+        test_cfg (dict): Config for testing. Default: None.
+        pretrained (str): Path to the pretrained models.
+        loss_pose (None): Deprecated arguments. Please use
+            ``loss_keypoint`` for heads instead.
+    """
+
+    def __init__(self,
+                 backbone,
+                 keypoint_head=None,
+                 train_cfg=None,
+                 test_cfg=None,
+                 pretrained=None,
+                 loss_pose=None):
+        super().__init__()
+        self.fp16_enabled = False
+
+        self.backbone = builder.build_backbone(backbone)
+
+        if keypoint_head is not None:
+            if 'loss_keypoint' not in keypoint_head and loss_pose is not None:
+                warnings.warn(
+                    '`loss_pose` for BottomUp is deprecated, '
+                    'use `loss_keypoint` for heads instead. See '
+                    'https://github.com/open-mmlab/mmpose/pull/382'
+                    ' for more information.', DeprecationWarning)
+                keypoint_head['loss_keypoint'] = loss_pose
+
+            self.keypoint_head = builder.build_head(keypoint_head)
+
+        self.train_cfg = train_cfg
+        self.test_cfg = test_cfg
+        self.use_udp = test_cfg.get('use_udp', False)
+        self.parser = HeatmapParser(self.test_cfg)
+        self.init_weights(pretrained=pretrained)
+
+    @property
+    def with_keypoint(self):
+        """Check if has keypoint_head."""
+        return hasattr(self, 'keypoint_head')
+
+    def init_weights(self, pretrained=None):
+        """Weight initialization for model."""
+        self.backbone.init_weights(pretrained)
+        if self.with_keypoint:
+            self.keypoint_head.init_weights()
+
+    @auto_fp16(apply_to=('img', ))
+    def forward(self,
+                img=None,
+                targets=None,
+                masks=None,
+                joints=None,
+                img_metas=None,
+                return_loss=True,
+                return_heatmap=False,
+                **kwargs):
+        """Calls either forward_train or forward_test depending on whether
+        return_loss is True.
+
+        Note:
+            - batch_size: N
+            - num_keypoints: K
+            - num_img_channel: C
+            - img_width: imgW
+            - img_height: imgH
+            - heatmaps weight: W
+            - heatmaps height: H
+            - max_num_people: M
+
+        Args:
+            img (torch.Tensor[N,C,imgH,imgW]): Input image.
+            targets (list(torch.Tensor[N,K,H,W])): Multi-scale target heatmaps.
+            masks (list(torch.Tensor[N,H,W])): Masks of multi-scale target
+                heatmaps
+            joints (list(torch.Tensor[N,M,K,2])): Joints of multi-scale target
+                heatmaps for ae loss
+            img_metas (dict): Information about val & test.
+                By default it includes:
+
+                - "image_file": image path
+                - "aug_data": input
+                - "test_scale_factor": test scale factor
+                - "base_size": base size of input
+                - "center": center of image
+                - "scale": scale of image
+                - "flip_index": flip index of keypoints
+            return loss (bool): ``return_loss=True`` for training,
+                ``return_loss=False`` for validation & test.
+            return_heatmap (bool) : Option to return heatmap.
+
+        Returns:
+            dict|tuple: if 'return_loss' is true, then return losses. \
+                Otherwise, return predicted poses, scores, image \
+                paths and heatmaps.
+        """
+
+        if return_loss:
+            return self.forward_train(img, targets, masks, joints, img_metas,
+                                      **kwargs)
+        return self.forward_test(
+            img, img_metas, return_heatmap=return_heatmap, **kwargs)
+
+    def forward_train(self, img, targets, masks, joints, img_metas, **kwargs):
+        """Forward the bottom-up model and calculate the loss.
+
+        Note:
+            batch_size: N
+            num_keypoints: K
+            num_img_channel: C
+            img_width: imgW
+            img_height: imgH
+            heatmaps weight: W
+            heatmaps height: H
+            max_num_people: M
+
+        Args:
+            img (torch.Tensor[N,C,imgH,imgW]): Input image.
+            targets (List(torch.Tensor[N,K,H,W])): Multi-scale target heatmaps.
+            masks (List(torch.Tensor[N,H,W])): Masks of multi-scale target
+                                              heatmaps
+            joints (List(torch.Tensor[N,M,K,2])): Joints of multi-scale target
+                                                 heatmaps for ae loss
+            img_metas (dict):Information about val&test
+                By default this includes:
+                - "image_file": image path
+                - "aug_data": input
+                - "test_scale_factor": test scale factor
+                - "base_size": base size of input
+                - "center": center of image
+                - "scale": scale of image
+                - "flip_index": flip index of keypoints
+
+        Returns:
+            dict: The total loss for bottom-up
+        """
+
+        output = self.backbone(img)
+
+        if self.with_keypoint:
+            output = self.keypoint_head(output)
+
+        # if return loss
+        losses = dict()
+        if self.with_keypoint:
+            keypoint_losses = self.keypoint_head.get_loss(
+                output, targets, masks, joints)
+            losses.update(keypoint_losses)
+
+        return losses
+
+    def forward_dummy(self, img):
+        """Used for computing network FLOPs.
+
+        See ``tools/get_flops.py``.
+
+        Args:
+            img (torch.Tensor): Input image.
+
+        Returns:
+            Tensor: Outputs.
+        """
+        output = self.backbone(img)
+        if self.with_keypoint:
+            output = self.keypoint_head(output)
+        return output
+
+    def forward_test(self, img, img_metas, return_heatmap=False, **kwargs):
+        """Inference the bottom-up model.
+
+        Note:
+            - Batchsize: N (currently support batchsize = 1)
+            - num_img_channel: C
+            - img_width: imgW
+            - img_height: imgH
+
+        Args:
+            flip_index (List(int)):
+            aug_data (List(Tensor[NxCximgHximgW])): Multi-scale image
+            test_scale_factor (List(float)): Multi-scale factor
+            base_size (Tuple(int)): Base size of image when scale is 1
+            center (np.ndarray): center of image
+            scale (np.ndarray): the scale of image
+        """
+        assert img.size(0) == 1
+        assert len(img_metas) == 1
+
+        img_metas = img_metas[0]
+
+        aug_data = img_metas['aug_data']
+
+        test_scale_factor = img_metas['test_scale_factor']
+        base_size = img_metas['base_size']
+        center = img_metas['center']
+        scale = img_metas['scale']
+
+        result = {}
+
+        scale_heatmaps_list = []
+        scale_tags_list = []
+
+        for idx, s in enumerate(sorted(test_scale_factor, reverse=True)):
+            image_resized = aug_data[idx].to(img.device)
+
+            features = self.backbone(image_resized)
+            if self.with_keypoint:
+                outputs = self.keypoint_head(features)
+
+            heatmaps, tags = split_ae_outputs(
+                outputs, self.test_cfg['num_joints'],
+                self.test_cfg['with_heatmaps'], self.test_cfg['with_ae'],
+                self.test_cfg.get('select_output_index', range(len(outputs))))
+
+            if self.test_cfg.get('flip_test', True):
+                # use flip test
+                features_flipped = self.backbone(
+                    torch.flip(image_resized, [3]))
+                if self.with_keypoint:
+                    outputs_flipped = self.keypoint_head(features_flipped)
+
+                heatmaps_flipped, tags_flipped = split_ae_outputs(
+                    outputs_flipped, self.test_cfg['num_joints'],
+                    self.test_cfg['with_heatmaps'], self.test_cfg['with_ae'],
+                    self.test_cfg.get('select_output_index',
+                                      range(len(outputs))))
+
+                heatmaps_flipped = flip_feature_maps(
+                    heatmaps_flipped, flip_index=img_metas['flip_index'])
+                if self.test_cfg['tag_per_joint']:
+                    tags_flipped = flip_feature_maps(
+                        tags_flipped, flip_index=img_metas['flip_index'])
+                else:
+                    tags_flipped = flip_feature_maps(
+                        tags_flipped, flip_index=None, flip_output=True)
+
+            else:
+                heatmaps_flipped = None
+                tags_flipped = None
+
+            aggregated_heatmaps = aggregate_stage_flip(
+                heatmaps,
+                heatmaps_flipped,
+                index=-1,
+                project2image=self.test_cfg['project2image'],
+                size_projected=base_size,
+                align_corners=self.test_cfg.get('align_corners', True),
+                aggregate_stage='average',
+                aggregate_flip='average')
+
+            aggregated_tags = aggregate_stage_flip(
+                tags,
+                tags_flipped,
+                index=-1,
+                project2image=self.test_cfg['project2image'],
+                size_projected=base_size,
+                align_corners=self.test_cfg.get('align_corners', True),
+                aggregate_stage='concat',
+                aggregate_flip='concat')
+
+            if s == 1 or len(test_scale_factor) == 1:
+                if isinstance(aggregated_tags, list):
+                    scale_tags_list.extend(aggregated_tags)
+                else:
+                    scale_tags_list.append(aggregated_tags)
+
+            if isinstance(aggregated_heatmaps, list):
+                scale_heatmaps_list.extend(aggregated_heatmaps)
+            else:
+                scale_heatmaps_list.append(aggregated_heatmaps)
+
+        aggregated_heatmaps = aggregate_scale(
+            scale_heatmaps_list,
+            align_corners=self.test_cfg.get('align_corners', True),
+            aggregate_scale='average')
+
+        aggregated_tags = aggregate_scale(
+            scale_tags_list,
+            align_corners=self.test_cfg.get('align_corners', True),
+            aggregate_scale='unsqueeze_concat')
+
+        heatmap_size = aggregated_heatmaps.shape[2:4]
+        tag_size = aggregated_tags.shape[2:4]
+        if heatmap_size != tag_size:
+            tmp = []
+            for idx in range(aggregated_tags.shape[-1]):
+                tmp.append(
+                    torch.nn.functional.interpolate(
+                        aggregated_tags[..., idx],
+                        size=heatmap_size,
+                        mode='bilinear',
+                        align_corners=self.test_cfg.get('align_corners',
+                                                        True)).unsqueeze(-1))
+            aggregated_tags = torch.cat(tmp, dim=-1)
+
+        # perform grouping
+        grouped, scores = self.parser.parse(aggregated_heatmaps,
+                                            aggregated_tags,
+                                            self.test_cfg['adjust'],
+                                            self.test_cfg['refine'])
+
+        preds = get_group_preds(
+            grouped,
+            center,
+            scale, [aggregated_heatmaps.size(3),
+                    aggregated_heatmaps.size(2)],
+            use_udp=self.use_udp)
+
+        image_paths = []
+        image_paths.append(img_metas['image_file'])
+
+        if return_heatmap:
+            output_heatmap = aggregated_heatmaps.detach().cpu().numpy()
+        else:
+            output_heatmap = None
+
+        result['preds'] = preds
+        result['scores'] = scores
+        result['image_paths'] = image_paths
+        result['output_heatmap'] = output_heatmap
+
+        return result
+
+    @deprecated_api_warning({'pose_limb_color': 'pose_link_color'},
+                            cls_name='AssociativeEmbedding')
+    def show_result(self,
+                    img,
+                    result,
+                    skeleton=None,
+                    kpt_score_thr=0.3,
+                    bbox_color=None,
+                    pose_kpt_color=None,
+                    pose_link_color=None,
+                    radius=4,
+                    thickness=1,
+                    font_scale=0.5,
+                    win_name='',
+                    show=False,
+                    show_keypoint_weight=False,
+                    wait_time=0,
+                    out_file=None):
+        """Draw `result` over `img`.
+
+        Args:
+            img (str or Tensor): The image to be displayed.
+            result (list[dict]): The results to draw over `img`
+                (bbox_result, pose_result).
+            skeleton (list[list]): The connection of keypoints.
+                skeleton is 0-based indexing.
+            kpt_score_thr (float, optional): Minimum score of keypoints
+                to be shown. Default: 0.3.
+            pose_kpt_color (np.array[Nx3]`): Color of N keypoints.
+                If None, do not draw keypoints.
+            pose_link_color (np.array[Mx3]): Color of M links.
+                If None, do not draw links.
+            radius (int): Radius of circles.
+            thickness (int): Thickness of lines.
+            font_scale (float): Font scales of texts.
+            win_name (str): The window name.
+            show (bool): Whether to show the image. Default: False.
+            show_keypoint_weight (bool): Whether to change the transparency
+                using the predicted confidence scores of keypoints.
+            wait_time (int): Value of waitKey param.
+                Default: 0.
+            out_file (str or None): The filename to write the image.
+                Default: None.
+
+        Returns:
+            Tensor: Visualized image only if not `show` or `out_file`
+        """
+        img = mmcv.imread(img)
+        img = img.copy()
+        img_h, img_w, _ = img.shape
+
+        pose_result = []
+        for res in result:
+            pose_result.append(res['keypoints'])
+
+        imshow_keypoints(img, pose_result, skeleton, kpt_score_thr,
+                         pose_kpt_color, pose_link_color, radius, thickness)
+
+        if show:
+            imshow(img, win_name, wait_time)
+
+        if out_file is not None:
+            imwrite(img, out_file)
+
+        return img

third_party/hamer/third-party/ViTPose/mmpose/models/detectors/base.py

@@ -0,0 +1,131 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from abc import ABCMeta, abstractmethod
+from collections import OrderedDict
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+
+
+class BasePose(nn.Module, metaclass=ABCMeta):
+    """Base class for pose detectors.
+
+    All recognizers should subclass it.
+    All subclass should overwrite:
+        Methods:`forward_train`, supporting to forward when training.
+        Methods:`forward_test`, supporting to forward when testing.
+
+    Args:
+        backbone (dict): Backbone modules to extract feature.
+        head (dict): Head modules to give output.
+        train_cfg (dict): Config for training. Default: None.
+        test_cfg (dict): Config for testing. Default: None.
+    """
+
+    @abstractmethod
+    def forward_train(self, img, img_metas, **kwargs):
+        """Defines the computation performed at training."""
+
+    @abstractmethod
+    def forward_test(self, img, img_metas, **kwargs):
+        """Defines the computation performed at testing."""
+
+    @abstractmethod
+    def forward(self, img, img_metas, return_loss=True, **kwargs):
+        """Forward function."""
+
+    @staticmethod
+    def _parse_losses(losses):
+        """Parse the raw outputs (losses) of the network.
+
+        Args:
+            losses (dict): Raw output of the network, which usually contain
+                losses and other necessary information.
+
+        Returns:
+            tuple[Tensor, dict]: (loss, log_vars), loss is the loss tensor \
+                which may be a weighted sum of all losses, log_vars \
+                contains all the variables to be sent to the logger.
+        """
+        log_vars = OrderedDict()
+        for loss_name, loss_value in losses.items():
+            if isinstance(loss_value, torch.Tensor):
+                log_vars[loss_name] = loss_value.mean()
+            elif isinstance(loss_value, float):
+                log_vars[loss_name] = loss_value
+            elif isinstance(loss_value, list):
+                log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value)
+            else:
+                raise TypeError(
+                    f'{loss_name} is not a tensor or list of tensors or float')
+
+        loss = sum(_value for _key, _value in log_vars.items()
+                   if 'loss' in _key)
+
+        log_vars['loss'] = loss
+        for loss_name, loss_value in log_vars.items():
+            # reduce loss when distributed training
+            if not isinstance(loss_value, float):
+                if dist.is_available() and dist.is_initialized():
+                    loss_value = loss_value.data.clone()
+                    dist.all_reduce(loss_value.div_(dist.get_world_size()))
+                log_vars[loss_name] = loss_value.item()
+            else:
+                log_vars[loss_name] = loss_value
+
+        return loss, log_vars
+
+    def train_step(self, data_batch, optimizer, **kwargs):
+        """The iteration step during training.
+
+        This method defines an iteration step during training, except for the
+        back propagation and optimizer updating, which are done in an optimizer
+        hook. Note that in some complicated cases or models, the whole process
+        including back propagation and optimizer updating is also defined in
+        this method, such as GAN.
+
+        Args:
+            data_batch (dict): The output of dataloader.
+            optimizer (:obj:`torch.optim.Optimizer` | dict): The optimizer of
+                runner is passed to ``train_step()``. This argument is unused
+                and reserved.
+
+        Returns:
+            dict: It should contain at least 3 keys: ``loss``, ``log_vars``,
+                ``num_samples``.
+                ``loss`` is a tensor for back propagation, which can be a
+                weighted sum of multiple losses.
+                ``log_vars`` contains all the variables to be sent to the
+                logger.
+                ``num_samples`` indicates the batch size (when the model is
+                DDP, it means the batch size on each GPU), which is used for
+                averaging the logs.
+        """
+        losses = self.forward(**data_batch)
+
+        loss, log_vars = self._parse_losses(losses)
+
+        outputs = dict(
+            loss=loss,
+            log_vars=log_vars,
+            num_samples=len(next(iter(data_batch.values()))))
+
+        return outputs
+
+    def val_step(self, data_batch, optimizer, **kwargs):
+        """The iteration step during validation.
+
+        This method shares the same signature as :func:`train_step`, but used
+        during val epochs. Note that the evaluation after training epochs is
+        not implemented with this method, but an evaluation hook.
+        """
+        results = self.forward(return_loss=False, **data_batch)
+
+        outputs = dict(results=results)
+
+        return outputs
+
+    @abstractmethod
+    def show_result(self, **kwargs):
+        """Visualize the results."""
+        raise NotImplementedError

third_party/hamer/third-party/ViTPose/mmpose/models/detectors/interhand_3d.py

@@ -0,0 +1,227 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import mmcv
+import numpy as np
+from mmcv.utils.misc import deprecated_api_warning
+
+from mmpose.core import imshow_keypoints, imshow_keypoints_3d
+from ..builder import POSENETS
+from .top_down import TopDown
+
+
+@POSENETS.register_module()
+class Interhand3D(TopDown):
+    """Top-down interhand 3D pose detector of paper ref: Gyeongsik Moon.
+
+    "InterHand2.6M: A Dataset and Baseline for 3D Interacting Hand Pose
+    Estimation from a Single RGB Image". A child class of TopDown detector.
+    """
+
+    def forward(self,
+                img,
+                target=None,
+                target_weight=None,
+                img_metas=None,
+                return_loss=True,
+                **kwargs):
+        """Calls either forward_train or forward_test depending on whether
+        return_loss=True. Note this setting will change the expected inputs.
+        When `return_loss=True`, img and img_meta are single-nested (i.e.
+        Tensor and List[dict]), and when `resturn_loss=False`, img and img_meta
+        should be double nested (i.e.  list[Tensor], list[list[dict]]), with
+        the outer list indicating test time augmentations.
+
+        Note:
+            - batch_size: N
+            - num_keypoints: K
+            - num_img_channel: C (Default: 3)
+            - img height: imgH
+            - img width: imgW
+            - heatmaps height: H
+            - heatmaps weight: W
+
+        Args:
+            img (torch.Tensor[NxCximgHximgW]): Input images.
+            target (list[torch.Tensor]): Target heatmaps, relative hand
+            root depth and hand type.
+            target_weight (list[torch.Tensor]): Weights for target
+            heatmaps, relative hand root depth and hand type.
+            img_metas (list(dict)): Information about data augmentation
+                By default this includes:
+
+                - "image_file: path to the image file
+                - "center": center of the bbox
+                - "scale": scale of the bbox
+                - "rotation": rotation of the bbox
+                - "bbox_score": score of bbox
+                - "heatmap3d_depth_bound": depth bound of hand keypoint 3D
+                    heatmap
+                - "root_depth_bound": depth bound of relative root depth 1D
+                    heatmap
+            return_loss (bool): Option to `return loss`. `return loss=True`
+                for training, `return loss=False` for validation & test.
+
+        Returns:
+            dict|tuple: if `return loss` is true, then return losses. \
+                Otherwise, return predicted poses, boxes, image paths, \
+                heatmaps, relative hand root depth and hand type.
+        """
+        if return_loss:
+            return self.forward_train(img, target, target_weight, img_metas,
+                                      **kwargs)
+        return self.forward_test(img, img_metas, **kwargs)
+
+    def forward_test(self, img, img_metas, **kwargs):
+        """Defines the computation performed at every call when testing."""
+        assert img.size(0) == len(img_metas)
+        batch_size, _, img_height, img_width = img.shape
+        if batch_size > 1:
+            assert 'bbox_id' in img_metas[0]
+
+        features = self.backbone(img)
+        if self.with_neck:
+            features = self.neck(features)
+        if self.with_keypoint:
+            output = self.keypoint_head.inference_model(
+                features, flip_pairs=None)
+
+        if self.test_cfg.get('flip_test', True):
+            img_flipped = img.flip(3)
+            features_flipped = self.backbone(img_flipped)
+            if self.with_neck:
+                features_flipped = self.neck(features_flipped)
+            if self.with_keypoint:
+                output_flipped = self.keypoint_head.inference_model(
+                    features_flipped, img_metas[0]['flip_pairs'])
+                output = [(out + out_flipped) * 0.5
+                          for out, out_flipped in zip(output, output_flipped)]
+
+        if self.with_keypoint:
+            result = self.keypoint_head.decode(
+                img_metas, output, img_size=[img_width, img_height])
+        else:
+            result = {}
+        return result
+
+    @deprecated_api_warning({'pose_limb_color': 'pose_link_color'},
+                            cls_name='Interhand3D')
+    def show_result(self,
+                    result,
+                    img=None,
+                    skeleton=None,
+                    kpt_score_thr=0.3,
+                    radius=8,
+                    bbox_color='green',
+                    thickness=2,
+                    pose_kpt_color=None,
+                    pose_link_color=None,
+                    vis_height=400,
+                    num_instances=-1,
+                    win_name='',
+                    show=False,
+                    wait_time=0,
+                    out_file=None):
+        """Visualize 3D pose estimation results.
+
+        Args:
+            result (list[dict]): The pose estimation results containing:
+
+                - "keypoints_3d" ([K,4]): 3D keypoints
+                - "keypoints" ([K,3] or [T,K,3]): Optional for visualizing
+                    2D inputs. If a sequence is given, only the last frame
+                    will be used for visualization
+                - "bbox" ([4,] or [T,4]): Optional for visualizing 2D inputs
+                - "title" (str): title for the subplot
+            img (str or Tensor): Optional. The image to visualize 2D inputs on.
+            skeleton (list of [idx_i,idx_j]): Skeleton described by a list of
+                links, each is a pair of joint indices.
+            kpt_score_thr (float, optional): Minimum score of keypoints
+                to be shown. Default: 0.3.
+            radius (int): Radius of circles.
+            bbox_color (str or tuple or :obj:`Color`): Color of bbox lines.
+            thickness (int): Thickness of lines.
+            pose_kpt_color (np.array[Nx3]`): Color of N keypoints.
+                If None, do not draw keypoints.
+            pose_link_color (np.array[Mx3]): Color of M limbs.
+                If None, do not draw limbs.
+            vis_height (int): The image height of the visualization. The width
+                will be N*vis_height depending on the number of visualized
+                items.
+            num_instances (int): Number of instances to be shown in 3D. If
+                smaller than 0, all the instances in the pose_result will be
+                shown. Otherwise, pad or truncate the pose_result to a length
+                of num_instances.
+            win_name (str): The window name.
+            show (bool): Whether to show the image. Default: False.
+            wait_time (int): Value of waitKey param.
+                Default: 0.
+            out_file (str or None): The filename to write the image.
+                Default: None.
+
+        Returns:
+            Tensor: Visualized img, only if not `show` or `out_file`.
+        """
+        if num_instances < 0:
+            assert len(result) > 0
+        result = sorted(result, key=lambda x: x.get('track_id', 0))
+
+        # draw image and 2d poses
+        if img is not None:
+            img = mmcv.imread(img)
+
+            bbox_result = []
+            pose_2d = []
+            for res in result:
+                if 'bbox' in res:
+                    bbox = np.array(res['bbox'])
+                    if bbox.ndim != 1:
+                        assert bbox.ndim == 2
+                        bbox = bbox[-1]  # Get bbox from the last frame
+                    bbox_result.append(bbox)
+                if 'keypoints' in res:
+                    kpts = np.array(res['keypoints'])
+                    if kpts.ndim != 2:
+                        assert kpts.ndim == 3
+                        kpts = kpts[-1]  # Get 2D keypoints from the last frame
+                    pose_2d.append(kpts)
+
+            if len(bbox_result) > 0:
+                bboxes = np.vstack(bbox_result)
+                mmcv.imshow_bboxes(
+                    img,
+                    bboxes,
+                    colors=bbox_color,
+                    top_k=-1,
+                    thickness=2,
+                    show=False)
+            if len(pose_2d) > 0:
+                imshow_keypoints(
+                    img,
+                    pose_2d,
+                    skeleton,
+                    kpt_score_thr=kpt_score_thr,
+                    pose_kpt_color=pose_kpt_color,
+                    pose_link_color=pose_link_color,
+                    radius=radius,
+                    thickness=thickness)
+            img = mmcv.imrescale(img, scale=vis_height / img.shape[0])
+
+        img_vis = imshow_keypoints_3d(
+            result,
+            img,
+            skeleton,
+            pose_kpt_color,
+            pose_link_color,
+            vis_height,
+            axis_limit=300,
+            axis_azimuth=-115,
+            axis_elev=15,
+            kpt_score_thr=kpt_score_thr,
+            num_instances=num_instances)
+
+        if show:
+            mmcv.visualization.imshow(img_vis, win_name, wait_time)
+
+        if out_file is not None:
+            mmcv.imwrite(img_vis, out_file)
+
+        return img_vis

third_party/hamer/third-party/ViTPose/mmpose/models/detectors/mesh.py

@@ -0,0 +1,438 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import cv2
+import mmcv
+import numpy as np
+import torch
+
+from mmpose.core.visualization.image import imshow_mesh_3d
+from mmpose.models.misc.discriminator import SMPLDiscriminator
+from .. import builder
+from ..builder import POSENETS
+from .base import BasePose
+
+
+def set_requires_grad(nets, requires_grad=False):
+    """Set requies_grad for all the networks.
+
+    Args:
+        nets (nn.Module | list[nn.Module]): A list of networks or a single
+            network.
+        requires_grad (bool): Whether the networks require gradients or not
+    """
+    if not isinstance(nets, list):
+        nets = [nets]
+    for net in nets:
+        if net is not None:
+            for param in net.parameters():
+                param.requires_grad = requires_grad
+
+
+@POSENETS.register_module()
+class ParametricMesh(BasePose):
+    """Model-based 3D human mesh detector. Take a single color image as input
+    and output 3D joints, SMPL parameters and camera parameters.
+
+    Args:
+        backbone (dict): Backbone modules to extract feature.
+        mesh_head (dict): Mesh head to process feature.
+        smpl (dict): Config for SMPL model.
+        disc (dict): Discriminator for SMPL parameters. Default: None.
+        loss_gan (dict): Config for adversarial loss. Default: None.
+        loss_mesh (dict): Config for mesh loss. Default: None.
+        train_cfg (dict): Config for training. Default: None.
+        test_cfg (dict): Config for testing. Default: None.
+        pretrained (str): Path to the pretrained models.
+    """
+
+    def __init__(self,
+                 backbone,
+                 mesh_head,
+                 smpl,
+                 disc=None,
+                 loss_gan=None,
+                 loss_mesh=None,
+                 train_cfg=None,
+                 test_cfg=None,
+                 pretrained=None):
+        super().__init__()
+
+        self.backbone = builder.build_backbone(backbone)
+        self.mesh_head = builder.build_head(mesh_head)
+        self.generator = torch.nn.Sequential(self.backbone, self.mesh_head)
+
+        self.smpl = builder.build_mesh_model(smpl)
+
+        self.with_gan = disc is not None and loss_gan is not None
+        if self.with_gan:
+            self.discriminator = SMPLDiscriminator(**disc)
+            self.loss_gan = builder.build_loss(loss_gan)
+        self.disc_step_count = 0
+
+        self.train_cfg = train_cfg
+        self.test_cfg = test_cfg
+
+        self.loss_mesh = builder.build_loss(loss_mesh)
+        self.init_weights(pretrained=pretrained)
+
+    def init_weights(self, pretrained=None):
+        """Weight initialization for model."""
+        self.backbone.init_weights(pretrained)
+        self.mesh_head.init_weights()
+        if self.with_gan:
+            self.discriminator.init_weights()
+
+    def train_step(self, data_batch, optimizer, **kwargs):
+        """Train step function.
+
+        In this function, the detector will finish the train step following
+        the pipeline:
+
+            1. get fake and real SMPL parameters
+            2. optimize discriminator (if have)
+            3. optimize generator
+
+        If `self.train_cfg.disc_step > 1`, the train step will contain multiple
+        iterations for optimizing discriminator with different input data and
+        only one iteration for optimizing generator after `disc_step`
+        iterations for discriminator.
+
+        Args:
+            data_batch (torch.Tensor): Batch of data as input.
+            optimizer (dict[torch.optim.Optimizer]): Dict with optimizers for
+                generator and discriminator (if have).
+
+        Returns:
+            outputs (dict): Dict with loss, information for logger,
+            the number of samples.
+        """
+
+        img = data_batch['img']
+        pred_smpl = self.generator(img)
+        pred_pose, pred_beta, pred_camera = pred_smpl
+
+        # optimize discriminator (if have)
+        if self.train_cfg['disc_step'] > 0 and self.with_gan:
+            set_requires_grad(self.discriminator, True)
+            fake_data = (pred_camera.detach(), pred_pose.detach(),
+                         pred_beta.detach())
+            mosh_theta = data_batch['mosh_theta']
+            real_data = (mosh_theta[:, :3], mosh_theta[:,
+                                                       3:75], mosh_theta[:,
+                                                                         75:])
+            fake_score = self.discriminator(fake_data)
+            real_score = self.discriminator(real_data)
+
+            disc_losses = {}
+            disc_losses['real_loss'] = self.loss_gan(
+                real_score, target_is_real=True, is_disc=True)
+            disc_losses['fake_loss'] = self.loss_gan(
+                fake_score, target_is_real=False, is_disc=True)
+            loss_disc, log_vars_d = self._parse_losses(disc_losses)
+
+            optimizer['discriminator'].zero_grad()
+            loss_disc.backward()
+            optimizer['discriminator'].step()
+            self.disc_step_count = \
+                (self.disc_step_count + 1) % self.train_cfg['disc_step']
+
+            if self.disc_step_count != 0:
+                outputs = dict(
+                    loss=loss_disc,
+                    log_vars=log_vars_d,
+                    num_samples=len(next(iter(data_batch.values()))))
+                return outputs
+
+        # optimize generator
+        pred_out = self.smpl(
+            betas=pred_beta,
+            body_pose=pred_pose[:, 1:],
+            global_orient=pred_pose[:, :1])
+        pred_vertices, pred_joints_3d = pred_out['vertices'], pred_out[
+            'joints']
+
+        gt_beta = data_batch['beta']
+        gt_pose = data_batch['pose']
+        gt_vertices = self.smpl(
+            betas=gt_beta,
+            body_pose=gt_pose[:, 3:],
+            global_orient=gt_pose[:, :3])['vertices']
+
+        pred = dict(
+            pose=pred_pose,
+            beta=pred_beta,
+            camera=pred_camera,
+            vertices=pred_vertices,
+            joints_3d=pred_joints_3d)
+
+        target = {
+            key: data_batch[key]
+            for key in [
+                'pose', 'beta', 'has_smpl', 'joints_3d', 'joints_2d',
+                'joints_3d_visible', 'joints_2d_visible'
+            ]
+        }
+        target['vertices'] = gt_vertices
+
+        losses = self.loss_mesh(pred, target)
+
+        if self.with_gan:
+            set_requires_grad(self.discriminator, False)
+            pred_theta = (pred_camera, pred_pose, pred_beta)
+            pred_score = self.discriminator(pred_theta)
+            loss_adv = self.loss_gan(
+                pred_score, target_is_real=True, is_disc=False)
+            losses['adv_loss'] = loss_adv
+
+        loss, log_vars = self._parse_losses(losses)
+        optimizer['generator'].zero_grad()
+        loss.backward()
+        optimizer['generator'].step()
+
+        outputs = dict(
+            loss=loss,
+            log_vars=log_vars,
+            num_samples=len(next(iter(data_batch.values()))))
+
+        return outputs
+
+    def forward_train(self, *args, **kwargs):
+        """Forward function for training.
+
+        For ParametricMesh, we do not use this interface.
+        """
+        raise NotImplementedError('This interface should not be used in '
+                                  'current training schedule. Please use '
+                                  '`train_step` for training.')
+
+    def val_step(self, data_batch, **kwargs):
+        """Forward function for evaluation.
+
+        Args:
+            data_batch (dict): Contain data for forward.
+
+        Returns:
+            dict: Contain the results from model.
+        """
+        output = self.forward_test(**data_batch, **kwargs)
+        return output
+
+    def forward_dummy(self, img):
+        """Used for computing network FLOPs.
+
+        See ``tools/get_flops.py``.
+
+        Args:
+            img (torch.Tensor): Input image.
+
+        Returns:
+            Tensor: Outputs.
+        """
+        output = self.generator(img)
+        return output
+
+    def forward_test(self,
+                     img,
+                     img_metas,
+                     return_vertices=False,
+                     return_faces=False,
+                     **kwargs):
+        """Defines the computation performed at every call when testing."""
+
+        pred_smpl = self.generator(img)
+        pred_pose, pred_beta, pred_camera = pred_smpl
+        pred_out = self.smpl(
+            betas=pred_beta,
+            body_pose=pred_pose[:, 1:],
+            global_orient=pred_pose[:, :1])
+        pred_vertices, pred_joints_3d = pred_out['vertices'], pred_out[
+            'joints']
+
+        all_preds = {}
+        all_preds['keypoints_3d'] = pred_joints_3d.detach().cpu().numpy()
+        all_preds['smpl_pose'] = pred_pose.detach().cpu().numpy()
+        all_preds['smpl_beta'] = pred_beta.detach().cpu().numpy()
+        all_preds['camera'] = pred_camera.detach().cpu().numpy()
+
+        if return_vertices:
+            all_preds['vertices'] = pred_vertices.detach().cpu().numpy()
+        if return_faces:
+            all_preds['faces'] = self.smpl.get_faces()
+
+        all_boxes = []
+        image_path = []
+        for img_meta in img_metas:
+            box = np.zeros(6, dtype=np.float32)
+            c = img_meta['center']
+            s = img_meta['scale']
+            if 'bbox_score' in img_metas:
+                score = np.array(img_metas['bbox_score']).reshape(-1)
+            else:
+                score = 1.0
+            box[0:2] = c
+            box[2:4] = s
+            box[4] = np.prod(s * 200.0, axis=0)
+            box[5] = score
+            all_boxes.append(box)
+            image_path.append(img_meta['image_file'])
+
+        all_preds['bboxes'] = np.stack(all_boxes, axis=0)
+        all_preds['image_path'] = image_path
+        return all_preds
+
+    def get_3d_joints_from_mesh(self, vertices):
+        """Get 3D joints from 3D mesh using predefined joints regressor."""
+        return torch.matmul(
+            self.joints_regressor.to(vertices.device), vertices)
+
+    def forward(self, img, img_metas=None, return_loss=False, **kwargs):
+        """Forward function.
+
+        Calls either forward_train or forward_test depending on whether
+        return_loss=True.
+
+        Note:
+            - batch_size: N
+            - num_img_channel: C (Default: 3)
+            - img height: imgH
+            - img width: imgW
+
+        Args:
+            img (torch.Tensor[N x C x imgH x imgW]): Input images.
+            img_metas (list(dict)): Information about data augmentation
+                By default this includes:
+
+                - "image_file: path to the image file
+                - "center": center of the bbox
+                - "scale": scale of the bbox
+                - "rotation": rotation of the bbox
+                - "bbox_score": score of bbox
+            return_loss (bool): Option to `return loss`. `return loss=True`
+                for training, `return loss=False` for validation & test.
+
+        Returns:
+            Return predicted 3D joints, SMPL parameters, boxes and image paths.
+        """
+
+        if return_loss:
+            return self.forward_train(img, img_metas, **kwargs)
+        return self.forward_test(img, img_metas, **kwargs)
+
+    def show_result(self,
+                    result,
+                    img,
+                    show=False,
+                    out_file=None,
+                    win_name='',
+                    wait_time=0,
+                    bbox_color='green',
+                    mesh_color=(76, 76, 204),
+                    **kwargs):
+        """Visualize 3D mesh estimation results.
+
+        Args:
+            result (list[dict]): The mesh estimation results containing:
+
+               - "bbox" (ndarray[4]): instance bounding bbox
+               - "center" (ndarray[2]): bbox center
+               - "scale" (ndarray[2]): bbox scale
+               - "keypoints_3d" (ndarray[K,3]): predicted 3D keypoints
+               - "camera" (ndarray[3]): camera parameters
+               - "vertices" (ndarray[V, 3]): predicted 3D vertices
+               - "faces" (ndarray[F, 3]): mesh faces
+            img (str or Tensor): Optional. The image to visualize 2D inputs on.
+            win_name (str): The window name.
+            show (bool): Whether to show the image. Default: False.
+            wait_time (int): Value of waitKey param. Default: 0.
+            out_file (str or None): The filename to write the image.
+                Default: None.
+            bbox_color (str or tuple or :obj:`Color`): Color of bbox lines.
+            mesh_color (str or tuple or :obj:`Color`): Color of mesh surface.
+
+        Returns:
+            ndarray: Visualized img, only if not `show` or `out_file`.
+        """
+
+        if img is not None:
+            img = mmcv.imread(img)
+
+        focal_length = self.loss_mesh.focal_length
+        H, W, C = img.shape
+        img_center = np.array([[0.5 * W], [0.5 * H]])
+
+        # show bounding boxes
+        bboxes = [res['bbox'] for res in result]
+        bboxes = np.vstack(bboxes)
+        mmcv.imshow_bboxes(
+            img, bboxes, colors=bbox_color, top_k=-1, thickness=2, show=False)
+
+        vertex_list = []
+        face_list = []
+        for res in result:
+            vertices = res['vertices']
+            faces = res['faces']
+            camera = res['camera']
+            camera_center = res['center']
+            scale = res['scale']
+
+            # predicted vertices are in root-relative space,
+            # we need to translate them to camera space.
+            translation = np.array([
+                camera[1], camera[2],
+                2 * focal_length / (scale[0] * 200.0 * camera[0] + 1e-9)
+            ])
+            mean_depth = vertices[:, -1].mean() + translation[-1]
+            translation[:2] += (camera_center -
+                                img_center[:, 0]) / focal_length * mean_depth
+            vertices += translation[None, :]
+
+            vertex_list.append(vertices)
+            face_list.append(faces)
+
+        # render from front view
+        img_vis = imshow_mesh_3d(
+            img,
+            vertex_list,
+            face_list,
+            img_center, [focal_length, focal_length],
+            colors=mesh_color)
+
+        # render from side view
+        # rotate mesh vertices
+        R = cv2.Rodrigues(np.array([0, np.radians(90.), 0]))[0]
+        rot_vertex_list = [np.dot(vert, R) for vert in vertex_list]
+
+        # get the 3D bbox containing all meshes
+        rot_vertices = np.concatenate(rot_vertex_list, axis=0)
+        min_corner = rot_vertices.min(0)
+        max_corner = rot_vertices.max(0)
+
+        center_3d = 0.5 * (min_corner + max_corner)
+        ratio = 0.8
+        bbox3d_size = max_corner - min_corner
+
+        # set appropriate translation to make all meshes appear in the image
+        z_x = bbox3d_size[0] * focal_length / (ratio * W) - min_corner[2]
+        z_y = bbox3d_size[1] * focal_length / (ratio * H) - min_corner[2]
+        z = max(z_x, z_y)
+        translation = -center_3d
+        translation[2] = z
+        translation = translation[None, :]
+        rot_vertex_list = [
+            rot_vert + translation for rot_vert in rot_vertex_list
+        ]
+
+        # render from side view
+        img_side = imshow_mesh_3d(
+            np.ones_like(img) * 255, rot_vertex_list, face_list, img_center,
+            [focal_length, focal_length])
+
+        # merger images from front view and side view
+        img_vis = np.concatenate([img_vis, img_side], axis=1)
+
+        if show:
+            mmcv.visualization.imshow(img_vis, win_name, wait_time)
+
+        if out_file is not None:
+            mmcv.imwrite(img_vis, out_file)
+
+        return img_vis

third_party/hamer/third-party/ViTPose/mmpose/models/detectors/multi_task.py

@@ -0,0 +1,187 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch.nn as nn
+
+from .. import builder
+from ..builder import POSENETS
+
+
+@POSENETS.register_module()
+class MultiTask(nn.Module):
+    """Multi-task detectors.
+
+    Args:
+        backbone (dict): Backbone modules to extract feature.
+        heads (list[dict]): heads to output predictions.
+        necks (list[dict] | None): necks to process feature.
+        head2neck (dict{int:int}): head index to neck index.
+        pretrained (str): Path to the pretrained models.
+    """
+
+    def __init__(self,
+                 backbone,
+                 heads,
+                 necks=None,
+                 head2neck=None,
+                 pretrained=None):
+        super().__init__()
+
+        self.backbone = builder.build_backbone(backbone)
+
+        if head2neck is None:
+            assert necks is None
+            head2neck = {}
+
+        self.head2neck = {}
+        for i in range(len(heads)):
+            self.head2neck[i] = head2neck[i] if i in head2neck else -1
+
+        self.necks = nn.ModuleList([])
+        if necks is not None:
+            for neck in necks:
+                self.necks.append(builder.build_neck(neck))
+        self.necks.append(nn.Identity())
+
+        self.heads = nn.ModuleList([])
+        assert heads is not None
+        for head in heads:
+            assert head is not None
+            self.heads.append(builder.build_head(head))
+
+        self.init_weights(pretrained=pretrained)
+
+    @property
+    def with_necks(self):
+        """Check if has keypoint_head."""
+        return hasattr(self, 'necks')
+
+    def init_weights(self, pretrained=None):
+        """Weight initialization for model."""
+        self.backbone.init_weights(pretrained)
+        if self.with_necks:
+            for neck in self.necks:
+                if hasattr(neck, 'init_weights'):
+                    neck.init_weights()
+
+        for head in self.heads:
+            if hasattr(head, 'init_weights'):
+                head.init_weights()
+
+    def forward(self,
+                img,
+                target=None,
+                target_weight=None,
+                img_metas=None,
+                return_loss=True,
+                **kwargs):
+        """Calls either forward_train or forward_test depending on whether
+        return_loss=True. Note this setting will change the expected inputs.
+        When `return_loss=True`, img and img_meta are single-nested (i.e.
+        Tensor and List[dict]), and when `resturn_loss=False`, img and img_meta
+        should be double nested (i.e.  List[Tensor], List[List[dict]]), with
+        the outer list indicating test time augmentations.
+
+        Note:
+            - batch_size: N
+            - num_keypoints: K
+            - num_img_channel: C (Default: 3)
+            - img height: imgH
+            - img weight: imgW
+            - heatmaps height: H
+            - heatmaps weight: W
+
+        Args:
+            img (torch.Tensor[N,C,imgH,imgW]): Input images.
+            target (list[torch.Tensor]): Targets.
+            target_weight (List[torch.Tensor]): Weights.
+            img_metas (list(dict)): Information about data augmentation
+                By default this includes:
+
+                - "image_file: path to the image file
+                - "center": center of the bbox
+                - "scale": scale of the bbox
+                - "rotation": rotation of the bbox
+                - "bbox_score": score of bbox
+            return_loss (bool): Option to `return loss`. `return loss=True`
+                for training, `return loss=False` for validation & test.
+
+        Returns:
+            dict|tuple: if `return loss` is true, then return losses. \
+                Otherwise, return predicted poses, boxes, image paths \
+                and heatmaps.
+        """
+        if return_loss:
+            return self.forward_train(img, target, target_weight, img_metas,
+                                      **kwargs)
+        return self.forward_test(img, img_metas, **kwargs)
+
+    def forward_train(self, img, target, target_weight, img_metas, **kwargs):
+        """Defines the computation performed at every call when training."""
+        features = self.backbone(img)
+        outputs = []
+
+        for head_id, head in enumerate(self.heads):
+            neck_id = self.head2neck[head_id]
+            outputs.append(head(self.necks[neck_id](features)))
+
+        # if return loss
+        losses = dict()
+
+        for head, output, gt, gt_weight in zip(self.heads, outputs, target,
+                                               target_weight):
+            loss = head.get_loss(output, gt, gt_weight)
+            assert len(set(losses.keys()).intersection(set(loss.keys()))) == 0
+            losses.update(loss)
+
+            if hasattr(head, 'get_accuracy'):
+                acc = head.get_accuracy(output, gt, gt_weight)
+                assert len(set(losses.keys()).intersection(set(
+                    acc.keys()))) == 0
+                losses.update(acc)
+
+        return losses
+
+    def forward_test(self, img, img_metas, **kwargs):
+        """Defines the computation performed at every call when testing."""
+        assert img.size(0) == len(img_metas)
+        batch_size, _, img_height, img_width = img.shape
+        if batch_size > 1:
+            assert 'bbox_id' in img_metas[0]
+
+        results = {}
+
+        features = self.backbone(img)
+        outputs = []
+
+        for head_id, head in enumerate(self.heads):
+            neck_id = self.head2neck[head_id]
+            if hasattr(head, 'inference_model'):
+                head_output = head.inference_model(
+                    self.necks[neck_id](features), flip_pairs=None)
+            else:
+                head_output = head(
+                    self.necks[neck_id](features)).detach().cpu().numpy()
+            outputs.append(head_output)
+
+        for head, output in zip(self.heads, outputs):
+            result = head.decode(
+                img_metas, output, img_size=[img_width, img_height])
+            results.update(result)
+        return results
+
+    def forward_dummy(self, img):
+        """Used for computing network FLOPs.
+
+        See ``tools/get_flops.py``.
+
+        Args:
+            img (torch.Tensor): Input image.
+
+        Returns:
+            list[Tensor]: Outputs.
+        """
+        features = self.backbone(img)
+        outputs = []
+        for head_id, head in enumerate(self.heads):
+            neck_id = self.head2neck[head_id]
+            outputs.append(head(self.necks[neck_id](features)))
+        return outputs

third_party/hamer/third-party/ViTPose/mmpose/models/detectors/multiview_pose.py

@@ -0,0 +1,889 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from mmcv.runner import load_checkpoint
+
+from mmpose.core.camera import SimpleCameraTorch
+from mmpose.core.post_processing.post_transforms import (
+    affine_transform_torch, get_affine_transform)
+from .. import builder
+from ..builder import POSENETS
+from .base import BasePose
+
+
+class ProjectLayer(nn.Module):
+
+    def __init__(self, image_size, heatmap_size):
+        """Project layer to get voxel feature. Adapted from
+        https://github.com/microsoft/voxelpose-
+        pytorch/blob/main/lib/models/project_layer.py.
+
+        Args:
+            image_size (int or list): input size of the 2D model
+            heatmap_size (int or list): output size of the 2D model
+        """
+        super(ProjectLayer, self).__init__()
+        self.image_size = image_size
+        self.heatmap_size = heatmap_size
+        if isinstance(self.image_size, int):
+            self.image_size = [self.image_size, self.image_size]
+        if isinstance(self.heatmap_size, int):
+            self.heatmap_size = [self.heatmap_size, self.heatmap_size]
+
+    def compute_grid(self, box_size, box_center, num_bins, device=None):
+        if isinstance(box_size, int) or isinstance(box_size, float):
+            box_size = [box_size, box_size, box_size]
+        if isinstance(num_bins, int):
+            num_bins = [num_bins, num_bins, num_bins]
+
+        grid_1D_x = torch.linspace(
+            -box_size[0] / 2, box_size[0] / 2, num_bins[0], device=device)
+        grid_1D_y = torch.linspace(
+            -box_size[1] / 2, box_size[1] / 2, num_bins[1], device=device)
+        grid_1D_z = torch.linspace(
+            -box_size[2] / 2, box_size[2] / 2, num_bins[2], device=device)
+        grid_x, grid_y, grid_z = torch.meshgrid(
+            grid_1D_x + box_center[0],
+            grid_1D_y + box_center[1],
+            grid_1D_z + box_center[2],
+        )
+        grid_x = grid_x.contiguous().view(-1, 1)
+        grid_y = grid_y.contiguous().view(-1, 1)
+        grid_z = grid_z.contiguous().view(-1, 1)
+        grid = torch.cat([grid_x, grid_y, grid_z], dim=1)
+
+        return grid
+
+    def get_voxel(self, feature_maps, meta, grid_size, grid_center, cube_size):
+        device = feature_maps[0].device
+        batch_size = feature_maps[0].shape[0]
+        num_channels = feature_maps[0].shape[1]
+        num_bins = cube_size[0] * cube_size[1] * cube_size[2]
+        n = len(feature_maps)
+        cubes = torch.zeros(
+            batch_size, num_channels, 1, num_bins, n, device=device)
+        w, h = self.heatmap_size
+        grids = torch.zeros(batch_size, num_bins, 3, device=device)
+        bounding = torch.zeros(batch_size, 1, 1, num_bins, n, device=device)
+        for i in range(batch_size):
+            if len(grid_center[0]) == 3 or grid_center[i][3] >= 0:
+                if len(grid_center) == 1:
+                    grid = self.compute_grid(
+                        grid_size, grid_center[0], cube_size, device=device)
+                else:
+                    grid = self.compute_grid(
+                        grid_size, grid_center[i], cube_size, device=device)
+                grids[i:i + 1] = grid
+                for c in range(n):
+                    center = meta[i]['center'][c]
+                    scale = meta[i]['scale'][c]
+
+                    width, height = center * 2
+                    trans = torch.as_tensor(
+                        get_affine_transform(center, scale / 200.0, 0,
+                                             self.image_size),
+                        dtype=torch.float,
+                        device=device)
+
+                    cam_param = meta[i]['camera'][c].copy()
+
+                    single_view_camera = SimpleCameraTorch(
+                        param=cam_param, device=device)
+                    xy = single_view_camera.world_to_pixel(grid)
+
+                    bounding[i, 0, 0, :, c] = (xy[:, 0] >= 0) & (
+                        xy[:, 1] >= 0) & (xy[:, 0] < width) & (
+                            xy[:, 1] < height)
+                    xy = torch.clamp(xy, -1.0, max(width, height))
+                    xy = affine_transform_torch(xy, trans)
+                    xy = xy * torch.tensor(
+                        [w, h], dtype=torch.float,
+                        device=device) / torch.tensor(
+                            self.image_size, dtype=torch.float, device=device)
+                    sample_grid = xy / torch.tensor([w - 1, h - 1],
+                                                    dtype=torch.float,
+                                                    device=device) * 2.0 - 1.0
+                    sample_grid = torch.clamp(
+                        sample_grid.view(1, 1, num_bins, 2), -1.1, 1.1)
+
+                    cubes[i:i + 1, :, :, :, c] += F.grid_sample(
+                        feature_maps[c][i:i + 1, :, :, :],
+                        sample_grid,
+                        align_corners=True)
+
+        cubes = torch.sum(
+            torch.mul(cubes, bounding), dim=-1) / (
+                torch.sum(bounding, dim=-1) + 1e-6)
+        cubes[cubes != cubes] = 0.0
+        cubes = cubes.clamp(0.0, 1.0)
+
+        cubes = cubes.view(batch_size, num_channels, cube_size[0],
+                           cube_size[1], cube_size[2])
+        return cubes, grids
+
+    def forward(self, feature_maps, meta, grid_size, grid_center, cube_size):
+        cubes, grids = self.get_voxel(feature_maps, meta, grid_size,
+                                      grid_center, cube_size)
+        return cubes, grids
+
+
+@POSENETS.register_module()
+class DetectAndRegress(BasePose):
+    """DetectAndRegress approach for multiview human pose detection.
+
+    Args:
+        backbone (ConfigDict): Dictionary to construct the 2D pose detector
+        human_detector (ConfigDict): dictionary to construct human detector
+        pose_regressor (ConfigDict): dictionary to construct pose regressor
+        train_cfg (ConfigDict): Config for training. Default: None.
+        test_cfg (ConfigDict): Config for testing. Default: None.
+        pretrained (str): Path to the pretrained 2D model. Default: None.
+        freeze_2d (bool): Whether to freeze the 2D model in training.
+            Default: True.
+    """
+
+    def __init__(self,
+                 backbone,
+                 human_detector,
+                 pose_regressor,
+                 train_cfg=None,
+                 test_cfg=None,
+                 pretrained=None,
+                 freeze_2d=True):
+        super(DetectAndRegress, self).__init__()
+        if backbone is not None:
+            self.backbone = builder.build_posenet(backbone)
+            if self.training and pretrained is not None:
+                load_checkpoint(self.backbone, pretrained)
+        else:
+            self.backbone = None
+
+        self.freeze_2d = freeze_2d
+        self.human_detector = builder.MODELS.build(human_detector)
+        self.pose_regressor = builder.MODELS.build(pose_regressor)
+
+        self.train_cfg = train_cfg
+        self.test_cfg = test_cfg
+
+    @staticmethod
+    def _freeze(model):
+        """Freeze parameters."""
+        model.eval()
+        for param in model.parameters():
+            param.requires_grad = False
+
+    def train(self, mode=True):
+        """Sets the module in training mode.
+        Args:
+            mode (bool): whether to set training mode (``True``)
+                or evaluation mode (``False``). Default: ``True``.
+
+        Returns:
+            Module: self
+        """
+        super().train(mode)
+        if mode and self.freeze_2d and self.backbone is not None:
+            self._freeze(self.backbone)
+
+        return self
+
+    def forward(self,
+                img=None,
+                img_metas=None,
+                return_loss=True,
+                targets=None,
+                masks=None,
+                targets_3d=None,
+                input_heatmaps=None,
+                **kwargs):
+        """
+        Note:
+            batch_size: N
+            num_keypoints: K
+            num_img_channel: C
+            img_width: imgW
+            img_height: imgH
+            feature_maps width: W
+            feature_maps height: H
+            volume_length: cubeL
+            volume_width: cubeW
+            volume_height: cubeH
+
+        Args:
+            img (list(torch.Tensor[NxCximgHximgW])):
+                Multi-camera input images to the 2D model.
+            img_metas (list(dict)):
+                Information about image, 3D groundtruth and camera parameters.
+            return_loss: Option to `return loss`. `return loss=True`
+                for training, `return loss=False` for validation & test.
+            targets (list(torch.Tensor[NxKxHxW])):
+                Multi-camera target feature_maps of the 2D model.
+            masks (list(torch.Tensor[NxHxW])):
+                Multi-camera masks of the input to the 2D model.
+            targets_3d (torch.Tensor[NxcubeLxcubeWxcubeH]):
+                Ground-truth 3D heatmap of human centers.
+            input_heatmaps (list(torch.Tensor[NxKxHxW])):
+                Multi-camera feature_maps when the 2D model is not available.
+                 Default: None.
+            **kwargs:
+
+        Returns:
+            dict: if 'return_loss' is true, then return losses.
+              Otherwise, return predicted poses, human centers and sample_id
+
+        """
+        if return_loss:
+            return self.forward_train(img, img_metas, targets, masks,
+                                      targets_3d, input_heatmaps)
+        else:
+            return self.forward_test(img, img_metas, input_heatmaps)
+
+    def train_step(self, data_batch, optimizer, **kwargs):
+        """The iteration step during training.
+
+        This method defines an iteration step during training, except for the
+        back propagation and optimizer updating, which are done in an optimizer
+        hook. Note that in some complicated cases or models, the whole process
+        including back propagation and optimizer updating is also defined in
+        this method, such as GAN.
+
+        Args:
+            data_batch (dict): The output of dataloader.
+            optimizer (:obj:`torch.optim.Optimizer` | dict): The optimizer of
+                runner is passed to ``train_step()``. This argument is unused
+                and reserved.
+
+        Returns:
+            dict: It should contain at least 3 keys: ``loss``, ``log_vars``,
+                ``num_samples``.
+                ``loss`` is a tensor for back propagation, which can be a
+                weighted sum of multiple losses.
+                ``log_vars`` contains all the variables to be sent to the
+                logger.
+                ``num_samples`` indicates the batch size (when the model is
+                DDP, it means the batch size on each GPU), which is used for
+                averaging the logs.
+        """
+        losses = self.forward(**data_batch)
+
+        loss, log_vars = self._parse_losses(losses)
+        if 'img' in data_batch:
+            batch_size = data_batch['img'][0].shape[0]
+        else:
+            assert 'input_heatmaps' in data_batch
+            batch_size = data_batch['input_heatmaps'][0][0].shape[0]
+
+        outputs = dict(loss=loss, log_vars=log_vars, num_samples=batch_size)
+
+        return outputs
+
+    def forward_train(self,
+                      img,
+                      img_metas,
+                      targets=None,
+                      masks=None,
+                      targets_3d=None,
+                      input_heatmaps=None):
+        """
+        Note:
+            batch_size: N
+            num_keypoints: K
+            num_img_channel: C
+            img_width: imgW
+            img_height: imgH
+            feature_maps width: W
+            feature_maps height: H
+            volume_length: cubeL
+            volume_width: cubeW
+            volume_height: cubeH
+
+        Args:
+            img (list(torch.Tensor[NxCximgHximgW])):
+                Multi-camera input images to the 2D model.
+            img_metas (list(dict)):
+                Information about image, 3D groundtruth and camera parameters.
+            targets (list(torch.Tensor[NxKxHxW])):
+                Multi-camera target feature_maps of the 2D model.
+            masks (list(torch.Tensor[NxHxW])):
+                Multi-camera masks of the input to the 2D model.
+            targets_3d (torch.Tensor[NxcubeLxcubeWxcubeH]):
+                Ground-truth 3D heatmap of human centers.
+            input_heatmaps (list(torch.Tensor[NxKxHxW])):
+                Multi-camera feature_maps when the 2D model is not available.
+                 Default: None.
+
+        Returns:
+            dict: losses.
+
+        """
+        if self.backbone is None:
+            assert input_heatmaps is not None
+            feature_maps = []
+            for input_heatmap in input_heatmaps:
+                feature_maps.append(input_heatmap[0])
+        else:
+            feature_maps = []
+            assert isinstance(img, list)
+            for img_ in img:
+                feature_maps.append(self.backbone.forward_dummy(img_)[0])
+
+        losses = dict()
+        human_candidates, human_loss = self.human_detector.forward_train(
+            None, img_metas, feature_maps, targets_3d, return_preds=True)
+        losses.update(human_loss)
+
+        pose_loss = self.pose_regressor(
+            None,
+            img_metas,
+            return_loss=True,
+            feature_maps=feature_maps,
+            human_candidates=human_candidates)
+        losses.update(pose_loss)
+
+        if not self.freeze_2d:
+            losses_2d = {}
+            heatmaps_tensor = torch.cat(feature_maps, dim=0)
+            targets_tensor = torch.cat(targets, dim=0)
+            masks_tensor = torch.cat(masks, dim=0)
+            losses_2d_ = self.backbone.get_loss(heatmaps_tensor,
+                                                targets_tensor, masks_tensor)
+            for k, v in losses_2d_.items():
+                losses_2d[k + '_2d'] = v
+            losses.update(losses_2d)
+
+        return losses
+
+    def forward_test(
+        self,
+        img,
+        img_metas,
+        input_heatmaps=None,
+    ):
+        """
+        Note:
+            batch_size: N
+            num_keypoints: K
+            num_img_channel: C
+            img_width: imgW
+            img_height: imgH
+            feature_maps width: W
+            feature_maps height: H
+            volume_length: cubeL
+            volume_width: cubeW
+            volume_height: cubeH
+
+        Args:
+            img (list(torch.Tensor[NxCximgHximgW])):
+                Multi-camera input images to the 2D model.
+            img_metas (list(dict)):
+                Information about image, 3D groundtruth and camera parameters.
+            input_heatmaps (list(torch.Tensor[NxKxHxW])):
+                Multi-camera feature_maps when the 2D model is not available.
+                 Default: None.
+
+        Returns:
+            dict: predicted poses, human centers and sample_id
+
+        """
+        if self.backbone is None:
+            assert input_heatmaps is not None
+            feature_maps = []
+            for input_heatmap in input_heatmaps:
+                feature_maps.append(input_heatmap[0])
+        else:
+            feature_maps = []
+            assert isinstance(img, list)
+            for img_ in img:
+                feature_maps.append(self.backbone.forward_dummy(img_)[0])
+
+        human_candidates = self.human_detector.forward_test(
+            None, img_metas, feature_maps)
+
+        human_poses = self.pose_regressor(
+            None,
+            img_metas,
+            return_loss=False,
+            feature_maps=feature_maps,
+            human_candidates=human_candidates)
+
+        result = {}
+        result['pose_3d'] = human_poses.cpu().numpy()
+        result['human_detection_3d'] = human_candidates.cpu().numpy()
+        result['sample_id'] = [img_meta['sample_id'] for img_meta in img_metas]
+
+        return result
+
+    def show_result(self, **kwargs):
+        """Visualize the results."""
+        raise NotImplementedError
+
+    def forward_dummy(self, img, input_heatmaps=None, num_candidates=5):
+        """Used for computing network FLOPs."""
+        if self.backbone is None:
+            assert input_heatmaps is not None
+            feature_maps = []
+            for input_heatmap in input_heatmaps:
+                feature_maps.append(input_heatmap[0])
+        else:
+            feature_maps = []
+            assert isinstance(img, list)
+            for img_ in img:
+                feature_maps.append(self.backbone.forward_dummy(img_)[0])
+
+        _ = self.human_detector.forward_dummy(feature_maps)
+
+        _ = self.pose_regressor.forward_dummy(feature_maps, num_candidates)
+
+
+@POSENETS.register_module()
+class VoxelSinglePose(BasePose):
+    """VoxelPose Please refer to the `paper <https://arxiv.org/abs/2004.06239>`
+    for details.
+
+    Args:
+        image_size (list): input size of the 2D model.
+        heatmap_size (list): output size of the 2D model.
+        sub_space_size (list): Size of the cuboid human proposal.
+        sub_cube_size (list): Size of the input volume to the pose net.
+        pose_net (ConfigDict): Dictionary to construct the pose net.
+        pose_head (ConfigDict): Dictionary to construct the pose head.
+        train_cfg (ConfigDict): Config for training. Default: None.
+        test_cfg (ConfigDict): Config for testing. Default: None.
+    """
+
+    def __init__(
+        self,
+        image_size,
+        heatmap_size,
+        sub_space_size,
+        sub_cube_size,
+        num_joints,
+        pose_net,
+        pose_head,
+        train_cfg=None,
+        test_cfg=None,
+    ):
+        super(VoxelSinglePose, self).__init__()
+        self.project_layer = ProjectLayer(image_size, heatmap_size)
+        self.pose_net = builder.build_backbone(pose_net)
+        self.pose_head = builder.build_head(pose_head)
+
+        self.sub_space_size = sub_space_size
+        self.sub_cube_size = sub_cube_size
+
+        self.num_joints = num_joints
+        self.train_cfg = train_cfg
+        self.test_cfg = test_cfg
+
+    def forward(self,
+                img,
+                img_metas,
+                return_loss=True,
+                feature_maps=None,
+                human_candidates=None,
+                **kwargs):
+        """
+        Note:
+            batch_size: N
+            num_keypoints: K
+            num_img_channel: C
+            img_width: imgW
+            img_height: imgH
+            feature_maps width: W
+            feature_maps height: H
+            volume_length: cubeL
+            volume_width: cubeW
+            volume_height: cubeH
+
+        Args:
+            img (list(torch.Tensor[NxCximgHximgW])):
+                Multi-camera input images to the 2D model.
+            feature_maps (list(torch.Tensor[NxCxHxW])):
+                Multi-camera input feature_maps.
+            img_metas (list(dict)):
+                Information about image, 3D groundtruth and camera parameters.
+            human_candidates (torch.Tensor[NxPx5]):
+                Human candidates.
+            return_loss: Option to `return loss`. `return loss=True`
+                for training, `return loss=False` for validation & test.
+
+        """
+        if return_loss:
+            return self.forward_train(img, img_metas, feature_maps,
+                                      human_candidates)
+        else:
+            return self.forward_test(img, img_metas, feature_maps,
+                                     human_candidates)
+
+    def forward_train(self,
+                      img,
+                      img_metas,
+                      feature_maps=None,
+                      human_candidates=None,
+                      return_preds=False,
+                      **kwargs):
+        """Defines the computation performed at training.
+        Note:
+            batch_size: N
+            num_keypoints: K
+            num_img_channel: C
+            img_width: imgW
+            img_height: imgH
+            feature_maps width: W
+            feature_maps height: H
+            volume_length: cubeL
+            volume_width: cubeW
+            volume_height: cubeH
+
+        Args:
+            img (list(torch.Tensor[NxCximgHximgW])):
+                Multi-camera input images to the 2D model.
+            feature_maps (list(torch.Tensor[NxCxHxW])):
+                Multi-camera input feature_maps.
+            img_metas (list(dict)):
+                Information about image, 3D groundtruth and camera parameters.
+            human_candidates (torch.Tensor[NxPx5]):
+                Human candidates.
+            return_preds (bool): Whether to return prediction results
+
+        Returns:
+            dict: losses.
+
+        """
+        batch_size, num_candidates, _ = human_candidates.shape
+        pred = human_candidates.new_zeros(batch_size, num_candidates,
+                                          self.num_joints, 5)
+        pred[:, :, :, 3:] = human_candidates[:, :, None, 3:]
+
+        device = feature_maps[0].device
+        gt_3d = torch.stack([
+            torch.tensor(img_meta['joints_3d'], device=device)
+            for img_meta in img_metas
+        ])
+        gt_3d_vis = torch.stack([
+            torch.tensor(img_meta['joints_3d_visible'], device=device)
+            for img_meta in img_metas
+        ])
+        valid_preds = []
+        valid_targets = []
+        valid_weights = []
+
+        for n in range(num_candidates):
+            index = pred[:, n, 0, 3] >= 0
+            num_valid = index.sum()
+            if num_valid > 0:
+                pose_input_cube, coordinates \
+                    = self.project_layer(feature_maps,
+                                         img_metas,
+                                         self.sub_space_size,
+                                         human_candidates[:, n, :3],
+                                         self.sub_cube_size)
+                pose_heatmaps_3d = self.pose_net(pose_input_cube)
+                pose_3d = self.pose_head(pose_heatmaps_3d[index],
+                                         coordinates[index])
+
+                pred[index, n, :, 0:3] = pose_3d.detach()
+                valid_targets.append(gt_3d[index, pred[index, n, 0, 3].long()])
+                valid_weights.append(gt_3d_vis[index, pred[index, n, 0,
+                                                           3].long(), :,
+                                               0:1].float())
+                valid_preds.append(pose_3d)
+
+        losses = dict()
+        if len(valid_preds) > 0:
+            valid_targets = torch.cat(valid_targets, dim=0)
+            valid_weights = torch.cat(valid_weights, dim=0)
+            valid_preds = torch.cat(valid_preds, dim=0)
+            losses.update(
+                self.pose_head.get_loss(valid_preds, valid_targets,
+                                        valid_weights))
+        else:
+            pose_input_cube = feature_maps[0].new_zeros(
+                batch_size, self.num_joints, *self.sub_cube_size)
+            coordinates = feature_maps[0].new_zeros(batch_size,
+                                                    *self.sub_cube_size,
+                                                    3).view(batch_size, -1, 3)
+            pseudo_targets = feature_maps[0].new_zeros(batch_size,
+                                                       self.num_joints, 3)
+            pseudo_weights = feature_maps[0].new_zeros(batch_size,
+                                                       self.num_joints, 1)
+            pose_heatmaps_3d = self.pose_net(pose_input_cube)
+            pose_3d = self.pose_head(pose_heatmaps_3d, coordinates)
+            losses.update(
+                self.pose_head.get_loss(pose_3d, pseudo_targets,
+                                        pseudo_weights))
+        if return_preds:
+            return pred, losses
+        else:
+            return losses
+
+    def forward_test(self,
+                     img,
+                     img_metas,
+                     feature_maps=None,
+                     human_candidates=None,
+                     **kwargs):
+        """Defines the computation performed at training.
+        Note:
+            batch_size: N
+            num_keypoints: K
+            num_img_channel: C
+            img_width: imgW
+            img_height: imgH
+            feature_maps width: W
+            feature_maps height: H
+            volume_length: cubeL
+            volume_width: cubeW
+            volume_height: cubeH
+
+        Args:
+            img (list(torch.Tensor[NxCximgHximgW])):
+                Multi-camera input images to the 2D model.
+            feature_maps (list(torch.Tensor[NxCxHxW])):
+                Multi-camera input feature_maps.
+            img_metas (list(dict)):
+                Information about image, 3D groundtruth and camera parameters.
+            human_candidates (torch.Tensor[NxPx5]):
+                Human candidates.
+
+        Returns:
+            dict: predicted poses, human centers and sample_id
+
+        """
+        batch_size, num_candidates, _ = human_candidates.shape
+        pred = human_candidates.new_zeros(batch_size, num_candidates,
+                                          self.num_joints, 5)
+        pred[:, :, :, 3:] = human_candidates[:, :, None, 3:]
+
+        for n in range(num_candidates):
+            index = pred[:, n, 0, 3] >= 0
+            num_valid = index.sum()
+            if num_valid > 0:
+                pose_input_cube, coordinates \
+                    = self.project_layer(feature_maps,
+                                         img_metas,
+                                         self.sub_space_size,
+                                         human_candidates[:, n, :3],
+                                         self.sub_cube_size)
+                pose_heatmaps_3d = self.pose_net(pose_input_cube)
+                pose_3d = self.pose_head(pose_heatmaps_3d[index],
+                                         coordinates[index])
+
+                pred[index, n, :, 0:3] = pose_3d.detach()
+
+        return pred
+
+    def show_result(self, **kwargs):
+        """Visualize the results."""
+        raise NotImplementedError
+
+    def forward_dummy(self, feature_maps, num_candidates=5):
+        """Used for computing network FLOPs."""
+        batch_size, num_channels = feature_maps[0].shape
+        pose_input_cube = feature_maps[0].new_zeros(batch_size, num_channels,
+                                                    *self.sub_cube_size)
+        for n in range(num_candidates):
+            _ = self.pose_net(pose_input_cube)
+
+
+@POSENETS.register_module()
+class VoxelCenterDetector(BasePose):
+    """Detect human center by 3D CNN on voxels.
+
+    Please refer to the
+    `paper <https://arxiv.org/abs/2004.06239>` for details.
+    Args:
+        image_size (list): input size of the 2D model.
+        heatmap_size (list): output size of the 2D model.
+        space_size (list): Size of the 3D space.
+        cube_size (list): Size of the input volume to the 3D CNN.
+        space_center (list): Coordinate of the center of the 3D space.
+        center_net (ConfigDict): Dictionary to construct the center net.
+        center_head (ConfigDict): Dictionary to construct the center head.
+        train_cfg (ConfigDict): Config for training. Default: None.
+        test_cfg (ConfigDict): Config for testing. Default: None.
+    """
+
+    def __init__(
+        self,
+        image_size,
+        heatmap_size,
+        space_size,
+        cube_size,
+        space_center,
+        center_net,
+        center_head,
+        train_cfg=None,
+        test_cfg=None,
+    ):
+        super(VoxelCenterDetector, self).__init__()
+        self.project_layer = ProjectLayer(image_size, heatmap_size)
+        self.center_net = builder.build_backbone(center_net)
+        self.center_head = builder.build_head(center_head)
+
+        self.space_size = space_size
+        self.cube_size = cube_size
+        self.space_center = space_center
+
+        self.train_cfg = train_cfg
+        self.test_cfg = test_cfg
+
+    def assign2gt(self, center_candidates, gt_centers, gt_num_persons):
+        """"Assign gt id to each valid human center candidate."""
+        det_centers = center_candidates[..., :3]
+        batch_size = center_candidates.shape[0]
+        cand_num = center_candidates.shape[1]
+        cand2gt = torch.zeros(batch_size, cand_num)
+
+        for i in range(batch_size):
+            cand = det_centers[i].view(cand_num, 1, -1)
+            gt = gt_centers[None, i, :gt_num_persons[i]]
+
+            dist = torch.sqrt(torch.sum((cand - gt)**2, dim=-1))
+            min_dist, min_gt = torch.min(dist, dim=-1)
+
+            cand2gt[i] = min_gt
+            cand2gt[i][min_dist > self.train_cfg['dist_threshold']] = -1.0
+
+        center_candidates[:, :, 3] = cand2gt
+
+        return center_candidates
+
+    def forward(self,
+                img,
+                img_metas,
+                return_loss=True,
+                feature_maps=None,
+                targets_3d=None):
+        """
+        Note:
+            batch_size: N
+            num_keypoints: K
+            num_img_channel: C
+            img_width: imgW
+            img_height: imgH
+            heatmaps width: W
+            heatmaps height: H
+        Args:
+            img (list(torch.Tensor[NxCximgHximgW])):
+                Multi-camera input images to the 2D model.
+            img_metas (list(dict)):
+                Information about image, 3D groundtruth and camera parameters.
+            return_loss: Option to `return loss`. `return loss=True`
+                for training, `return loss=False` for validation & test.
+            targets_3d (torch.Tensor[NxcubeLxcubeWxcubeH]):
+                Ground-truth 3D heatmap of human centers.
+            feature_maps (list(torch.Tensor[NxKxHxW])):
+                Multi-camera feature_maps.
+        Returns:
+            dict: if 'return_loss' is true, then return losses.
+                Otherwise, return predicted poses
+        """
+        if return_loss:
+            return self.forward_train(img, img_metas, feature_maps, targets_3d)
+        else:
+            return self.forward_test(img, img_metas, feature_maps)
+
+    def forward_train(self,
+                      img,
+                      img_metas,
+                      feature_maps=None,
+                      targets_3d=None,
+                      return_preds=False):
+        """
+        Note:
+            batch_size: N
+            num_keypoints: K
+            num_img_channel: C
+            img_width: imgW
+            img_height: imgH
+            heatmaps width: W
+            heatmaps height: H
+        Args:
+            img (list(torch.Tensor[NxCximgHximgW])):
+                Multi-camera input images to the 2D model.
+            img_metas (list(dict)):
+                Information about image, 3D groundtruth and camera parameters.
+            targets_3d (torch.Tensor[NxcubeLxcubeWxcubeH]):
+                Ground-truth 3D heatmap of human centers.
+            feature_maps (list(torch.Tensor[NxKxHxW])):
+                Multi-camera feature_maps.
+            return_preds (bool): Whether to return prediction results
+        Returns:
+            dict: if 'return_pred' is true, then return losses
+                and human centers. Otherwise, return losses only
+        """
+        initial_cubes, _ = self.project_layer(feature_maps, img_metas,
+                                              self.space_size,
+                                              [self.space_center],
+                                              self.cube_size)
+        center_heatmaps_3d = self.center_net(initial_cubes)
+        center_heatmaps_3d = center_heatmaps_3d.squeeze(1)
+        center_candidates = self.center_head(center_heatmaps_3d)
+
+        device = center_candidates.device
+
+        gt_centers = torch.stack([
+            torch.tensor(img_meta['roots_3d'], device=device)
+            for img_meta in img_metas
+        ])
+        gt_num_persons = torch.stack([
+            torch.tensor(img_meta['num_persons'], device=device)
+            for img_meta in img_metas
+        ])
+        center_candidates = self.assign2gt(center_candidates, gt_centers,
+                                           gt_num_persons)
+
+        losses = dict()
+        losses.update(
+            self.center_head.get_loss(center_heatmaps_3d, targets_3d))
+
+        if return_preds:
+            return center_candidates, losses
+        else:
+            return losses
+
+    def forward_test(self, img, img_metas, feature_maps=None):
+        """
+        Note:
+            batch_size: N
+            num_keypoints: K
+            num_img_channel: C
+            img_width: imgW
+            img_height: imgH
+            heatmaps width: W
+            heatmaps height: H
+        Args:
+            img (list(torch.Tensor[NxCximgHximgW])):
+                Multi-camera input images to the 2D model.
+            img_metas (list(dict)):
+                Information about image, 3D groundtruth and camera parameters.
+            feature_maps (list(torch.Tensor[NxKxHxW])):
+                Multi-camera feature_maps.
+        Returns:
+            human centers
+        """
+        initial_cubes, _ = self.project_layer(feature_maps, img_metas,
+                                              self.space_size,
+                                              [self.space_center],
+                                              self.cube_size)
+        center_heatmaps_3d = self.center_net(initial_cubes)
+        center_heatmaps_3d = center_heatmaps_3d.squeeze(1)
+        center_candidates = self.center_head(center_heatmaps_3d)
+        center_candidates[..., 3] = \
+            (center_candidates[..., 4] >
+             self.test_cfg['center_threshold']).float() - 1.0
+
+        return center_candidates
+
+    def show_result(self, **kwargs):
+        """Visualize the results."""
+        raise NotImplementedError
+
+    def forward_dummy(self, feature_maps):
+        """Used for computing network FLOPs."""
+        batch_size, num_channels, _, _ = feature_maps[0].shape
+        initial_cubes = feature_maps[0].new_zeros(batch_size, num_channels,
+                                                  *self.cube_size)
+        _ = self.center_net(initial_cubes)

third_party/hamer/third-party/ViTPose/mmpose/models/detectors/pose_lifter.py

@@ -0,0 +1,392 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import warnings
+
+import mmcv
+import numpy as np
+from mmcv.utils.misc import deprecated_api_warning
+
+from mmpose.core import imshow_bboxes, imshow_keypoints, imshow_keypoints_3d
+from .. import builder
+from ..builder import POSENETS
+from .base import BasePose
+
+try:
+    from mmcv.runner import auto_fp16
+except ImportError:
+    warnings.warn('auto_fp16 from mmpose will be deprecated from v0.15.0'
+                  'Please install mmcv>=1.1.4')
+    from mmpose.core import auto_fp16
+
+
+@POSENETS.register_module()
+class PoseLifter(BasePose):
+    """Pose lifter that lifts 2D pose to 3D pose.
+
+    The basic model is a pose model that predicts root-relative pose. If
+    traj_head is not None, a trajectory model that predicts absolute root joint
+    position is also built.
+
+    Args:
+        backbone (dict): Config for the backbone of pose model.
+        neck (dict|None): Config for the neck of pose model.
+        keypoint_head (dict|None): Config for the head of pose model.
+        traj_backbone (dict|None): Config for the backbone of trajectory model.
+            If traj_backbone is None and traj_head is not None, trajectory
+            model will share backbone with pose model.
+        traj_neck (dict|None): Config for the neck of trajectory model.
+        traj_head (dict|None): Config for the head of trajectory model.
+        loss_semi (dict|None): Config for semi-supervision loss.
+        train_cfg (dict|None): Config for keypoint head during training.
+        test_cfg (dict|None): Config for keypoint head during testing.
+        pretrained (str|None): Path to pretrained weights.
+    """
+
+    def __init__(self,
+                 backbone,
+                 neck=None,
+                 keypoint_head=None,
+                 traj_backbone=None,
+                 traj_neck=None,
+                 traj_head=None,
+                 loss_semi=None,
+                 train_cfg=None,
+                 test_cfg=None,
+                 pretrained=None):
+        super().__init__()
+        self.fp16_enabled = False
+
+        self.train_cfg = train_cfg
+        self.test_cfg = test_cfg
+
+        # pose model
+        self.backbone = builder.build_backbone(backbone)
+
+        if neck is not None:
+            self.neck = builder.build_neck(neck)
+
+        if keypoint_head is not None:
+            keypoint_head['train_cfg'] = train_cfg
+            keypoint_head['test_cfg'] = test_cfg
+            self.keypoint_head = builder.build_head(keypoint_head)
+
+        # trajectory model
+        if traj_head is not None:
+            self.traj_head = builder.build_head(traj_head)
+
+            if traj_backbone is not None:
+                self.traj_backbone = builder.build_backbone(traj_backbone)
+            else:
+                self.traj_backbone = self.backbone
+
+            if traj_neck is not None:
+                self.traj_neck = builder.build_neck(traj_neck)
+
+        # semi-supervised learning
+        self.semi = loss_semi is not None
+        if self.semi:
+            assert keypoint_head is not None and traj_head is not None
+            self.loss_semi = builder.build_loss(loss_semi)
+
+        self.init_weights(pretrained=pretrained)
+
+    @property
+    def with_neck(self):
+        """Check if has keypoint_neck."""
+        return hasattr(self, 'neck')
+
+    @property
+    def with_keypoint(self):
+        """Check if has keypoint_head."""
+        return hasattr(self, 'keypoint_head')
+
+    @property
+    def with_traj_backbone(self):
+        """Check if has trajectory_backbone."""
+        return hasattr(self, 'traj_backbone')
+
+    @property
+    def with_traj_neck(self):
+        """Check if has trajectory_neck."""
+        return hasattr(self, 'traj_neck')
+
+    @property
+    def with_traj(self):
+        """Check if has trajectory_head."""
+        return hasattr(self, 'traj_head')
+
+    @property
+    def causal(self):
+        if hasattr(self.backbone, 'causal'):
+            return self.backbone.causal
+        else:
+            raise AttributeError('A PoseLifter\'s backbone should have '
+                                 'the bool attribute "causal" to indicate if'
+                                 'it performs causal inference.')
+
+    def init_weights(self, pretrained=None):
+        """Weight initialization for model."""
+        self.backbone.init_weights(pretrained)
+        if self.with_neck:
+            self.neck.init_weights()
+        if self.with_keypoint:
+            self.keypoint_head.init_weights()
+        if self.with_traj_backbone:
+            self.traj_backbone.init_weights(pretrained)
+        if self.with_traj_neck:
+            self.traj_neck.init_weights()
+        if self.with_traj:
+            self.traj_head.init_weights()
+
+    @auto_fp16(apply_to=('input', ))
+    def forward(self,
+                input,
+                target=None,
+                target_weight=None,
+                metas=None,
+                return_loss=True,
+                **kwargs):
+        """Calls either forward_train or forward_test depending on whether
+        return_loss=True.
+
+        Note:
+            - batch_size: N
+            - num_input_keypoints: Ki
+            - input_keypoint_dim: Ci
+            - input_sequence_len: Ti
+            - num_output_keypoints: Ko
+            - output_keypoint_dim: Co
+            - input_sequence_len: To
+
+        Args:
+            input (torch.Tensor[NxKixCixTi]): Input keypoint coordinates.
+            target (torch.Tensor[NxKoxCoxTo]): Output keypoint coordinates.
+                Defaults to None.
+            target_weight (torch.Tensor[NxKox1]): Weights across different
+                joint types. Defaults to None.
+            metas (list(dict)): Information about data augmentation
+            return_loss (bool): Option to `return loss`. `return loss=True`
+                for training, `return loss=False` for validation & test.
+
+        Returns:
+            dict|Tensor: If `reutrn_loss` is true, return losses. \
+                Otherwise return predicted poses.
+        """
+        if return_loss:
+            return self.forward_train(input, target, target_weight, metas,
+                                      **kwargs)
+        else:
+            return self.forward_test(input, metas, **kwargs)
+
+    def forward_train(self, input, target, target_weight, metas, **kwargs):
+        """Defines the computation performed at every call when training."""
+        assert input.size(0) == len(metas)
+
+        # supervised learning
+        # pose model
+        features = self.backbone(input)
+        if self.with_neck:
+            features = self.neck(features)
+        if self.with_keypoint:
+            output = self.keypoint_head(features)
+
+        losses = dict()
+        if self.with_keypoint:
+            keypoint_losses = self.keypoint_head.get_loss(
+                output, target, target_weight)
+            keypoint_accuracy = self.keypoint_head.get_accuracy(
+                output, target, target_weight, metas)
+            losses.update(keypoint_losses)
+            losses.update(keypoint_accuracy)
+
+        # trajectory model
+        if self.with_traj:
+            traj_features = self.traj_backbone(input)
+            if self.with_traj_neck:
+                traj_features = self.traj_neck(traj_features)
+            traj_output = self.traj_head(traj_features)
+
+            traj_losses = self.traj_head.get_loss(traj_output,
+                                                  kwargs['traj_target'], None)
+            losses.update(traj_losses)
+
+        # semi-supervised learning
+        if self.semi:
+            ul_input = kwargs['unlabeled_input']
+            ul_features = self.backbone(ul_input)
+            if self.with_neck:
+                ul_features = self.neck(ul_features)
+            ul_output = self.keypoint_head(ul_features)
+
+            ul_traj_features = self.traj_backbone(ul_input)
+            if self.with_traj_neck:
+                ul_traj_features = self.traj_neck(ul_traj_features)
+            ul_traj_output = self.traj_head(ul_traj_features)
+
+            output_semi = dict(
+                labeled_pose=output,
+                unlabeled_pose=ul_output,
+                unlabeled_traj=ul_traj_output)
+            target_semi = dict(
+                unlabeled_target_2d=kwargs['unlabeled_target_2d'],
+                intrinsics=kwargs['intrinsics'])
+
+            semi_losses = self.loss_semi(output_semi, target_semi)
+            losses.update(semi_losses)
+
+        return losses
+
+    def forward_test(self, input, metas, **kwargs):
+        """Defines the computation performed at every call when training."""
+        assert input.size(0) == len(metas)
+
+        results = {}
+
+        features = self.backbone(input)
+        if self.with_neck:
+            features = self.neck(features)
+        if self.with_keypoint:
+            output = self.keypoint_head.inference_model(features)
+            keypoint_result = self.keypoint_head.decode(metas, output)
+            results.update(keypoint_result)
+
+        if self.with_traj:
+            traj_features = self.traj_backbone(input)
+            if self.with_traj_neck:
+                traj_features = self.traj_neck(traj_features)
+            traj_output = self.traj_head.inference_model(traj_features)
+            results['traj_preds'] = traj_output
+
+        return results
+
+    def forward_dummy(self, input):
+        """Used for computing network FLOPs. See ``tools/get_flops.py``.
+
+        Args:
+            input (torch.Tensor): Input pose
+
+        Returns:
+            Tensor: Model output
+        """
+        output = self.backbone(input)
+        if self.with_neck:
+            output = self.neck(output)
+        if self.with_keypoint:
+            output = self.keypoint_head(output)
+
+        if self.with_traj:
+            traj_features = self.traj_backbone(input)
+            if self.with_neck:
+                traj_features = self.traj_neck(traj_features)
+            traj_output = self.traj_head(traj_features)
+            output = output + traj_output
+
+        return output
+
+    @deprecated_api_warning({'pose_limb_color': 'pose_link_color'},
+                            cls_name='PoseLifter')
+    def show_result(self,
+                    result,
+                    img=None,
+                    skeleton=None,
+                    pose_kpt_color=None,
+                    pose_link_color=None,
+                    radius=8,
+                    thickness=2,
+                    vis_height=400,
+                    num_instances=-1,
+                    win_name='',
+                    show=False,
+                    wait_time=0,
+                    out_file=None):
+        """Visualize 3D pose estimation results.
+
+        Args:
+            result (list[dict]): The pose estimation results containing:
+
+                - "keypoints_3d" ([K,4]): 3D keypoints
+                - "keypoints" ([K,3] or [T,K,3]): Optional for visualizing
+                    2D inputs. If a sequence is given, only the last frame
+                    will be used for visualization
+                - "bbox" ([4,] or [T,4]): Optional for visualizing 2D inputs
+                - "title" (str): title for the subplot
+            img (str or Tensor): Optional. The image to visualize 2D inputs on.
+            skeleton (list of [idx_i,idx_j]): Skeleton described by a list of
+                links, each is a pair of joint indices.
+            pose_kpt_color (np.array[Nx3]`): Color of N keypoints.
+                If None, do not draw keypoints.
+            pose_link_color (np.array[Mx3]): Color of M links.
+                If None, do not draw links.
+            radius (int): Radius of circles.
+            thickness (int): Thickness of lines.
+            vis_height (int): The image height of the visualization. The width
+                will be N*vis_height depending on the number of visualized
+                items.
+            win_name (str): The window name.
+            wait_time (int): Value of waitKey param.
+                Default: 0.
+            out_file (str or None): The filename to write the image.
+                Default: None.
+
+        Returns:
+            Tensor: Visualized img, only if not `show` or `out_file`.
+        """
+        if num_instances < 0:
+            assert len(result) > 0
+        result = sorted(result, key=lambda x: x.get('track_id', 1e4))
+
+        # draw image and input 2d poses
+        if img is not None:
+            img = mmcv.imread(img)
+
+            bbox_result = []
+            pose_input_2d = []
+            for res in result:
+                if 'bbox' in res:
+                    bbox = np.array(res['bbox'])
+                    if bbox.ndim != 1:
+                        assert bbox.ndim == 2
+                        bbox = bbox[-1]  # Get bbox from the last frame
+                    bbox_result.append(bbox)
+                if 'keypoints' in res:
+                    kpts = np.array(res['keypoints'])
+                    if kpts.ndim != 2:
+                        assert kpts.ndim == 3
+                        kpts = kpts[-1]  # Get 2D keypoints from the last frame
+                    pose_input_2d.append(kpts)
+
+            if len(bbox_result) > 0:
+                bboxes = np.vstack(bbox_result)
+                imshow_bboxes(
+                    img,
+                    bboxes,
+                    colors='green',
+                    thickness=thickness,
+                    show=False)
+            if len(pose_input_2d) > 0:
+                imshow_keypoints(
+                    img,
+                    pose_input_2d,
+                    skeleton,
+                    kpt_score_thr=0.3,
+                    pose_kpt_color=pose_kpt_color,
+                    pose_link_color=pose_link_color,
+                    radius=radius,
+                    thickness=thickness)
+            img = mmcv.imrescale(img, scale=vis_height / img.shape[0])
+
+        img_vis = imshow_keypoints_3d(
+            result,
+            img,
+            skeleton,
+            pose_kpt_color,
+            pose_link_color,
+            vis_height,
+            num_instances=num_instances)
+
+        if show:
+            mmcv.visualization.imshow(img_vis, win_name, wait_time)
+
+        if out_file is not None:
+            mmcv.imwrite(img_vis, out_file)
+
+        return img_vis

third_party/hamer/third-party/ViTPose/mmpose/models/detectors/posewarper.py

@@ -0,0 +1,244 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import warnings
+
+import numpy as np
+import torch
+
+from ..builder import POSENETS
+from .top_down import TopDown
+
+try:
+    from mmcv.runner import auto_fp16
+except ImportError:
+    warnings.warn('auto_fp16 from mmpose will be deprecated from v0.15.0'
+                  'Please install mmcv>=1.1.4')
+    from mmpose.core import auto_fp16
+
+
+@POSENETS.register_module()
+class PoseWarper(TopDown):
+    """Top-down pose detectors for multi-frame settings for video inputs.
+
+    `"Learning temporal pose estimation from sparsely-labeled videos"
+    <https://arxiv.org/abs/1906.04016>`_.
+
+    A child class of TopDown detector. The main difference between PoseWarper
+    and TopDown lies in that the former takes a list of tensors as input image
+    while the latter takes a single tensor as input image in forward method.
+
+    Args:
+        backbone (dict): Backbone modules to extract features.
+        neck (dict): intermediate modules to transform features.
+        keypoint_head (dict): Keypoint head to process feature.
+        train_cfg (dict): Config for training. Default: None.
+        test_cfg (dict): Config for testing. Default: None.
+        pretrained (str): Path to the pretrained models.
+        loss_pose (None): Deprecated arguments. Please use
+            `loss_keypoint` for heads instead.
+        concat_tensors (bool): Whether to concat the tensors on the batch dim,
+            which can speed up, Default: True
+    """
+
+    def __init__(self,
+                 backbone,
+                 neck=None,
+                 keypoint_head=None,
+                 train_cfg=None,
+                 test_cfg=None,
+                 pretrained=None,
+                 loss_pose=None,
+                 concat_tensors=True):
+        super().__init__(
+            backbone=backbone,
+            neck=neck,
+            keypoint_head=keypoint_head,
+            train_cfg=train_cfg,
+            test_cfg=test_cfg,
+            pretrained=pretrained,
+            loss_pose=loss_pose)
+        self.concat_tensors = concat_tensors
+
+    @auto_fp16(apply_to=('img', ))
+    def forward(self,
+                img,
+                target=None,
+                target_weight=None,
+                img_metas=None,
+                return_loss=True,
+                return_heatmap=False,
+                **kwargs):
+        """Calls either forward_train or forward_test depending on whether
+        return_loss=True. Note this setting will change the expected inputs.
+        When `return_loss=True`, img and img_meta are single-nested (i.e.
+        Tensor and List[dict]), and when `resturn_loss=False`, img and img_meta
+        should be double nested (i.e.  List[Tensor], List[List[dict]]), with
+        the outer list indicating test time augmentations.
+
+        Note:
+            - number of frames: F
+            - batch_size: N
+            - num_keypoints: K
+            - num_img_channel: C (Default: 3)
+            - img height: imgH
+            - img width: imgW
+            - heatmaps height: H
+            - heatmaps weight: W
+
+        Args:
+            imgs (list[F,torch.Tensor[N,C,imgH,imgW]]): multiple input frames
+            target (torch.Tensor[N,K,H,W]): Target heatmaps for one frame.
+            target_weight (torch.Tensor[N,K,1]): Weights across
+                different joint types.
+            img_metas (list(dict)): Information about data augmentation
+                By default this includes:
+
+                - "image_file: paths to multiple video frames
+                - "center": center of the bbox
+                - "scale": scale of the bbox
+                - "rotation": rotation of the bbox
+                - "bbox_score": score of bbox
+            return_loss (bool): Option to `return loss`. `return loss=True`
+                for training, `return loss=False` for validation & test.
+            return_heatmap (bool) : Option to return heatmap.
+
+        Returns:
+            dict|tuple: if `return loss` is true, then return losses. \
+                Otherwise, return predicted poses, boxes, image paths \
+                and heatmaps.
+        """
+        if return_loss:
+            return self.forward_train(img, target, target_weight, img_metas,
+                                      **kwargs)
+        return self.forward_test(
+            img, img_metas, return_heatmap=return_heatmap, **kwargs)
+
+    def forward_train(self, imgs, target, target_weight, img_metas, **kwargs):
+        """Defines the computation performed at every call when training."""
+        # imgs (list[Fxtorch.Tensor[NxCximgHximgW]]): multiple input frames
+        assert imgs[0].size(0) == len(img_metas)
+        num_frames = len(imgs)
+        frame_weight = img_metas[0]['frame_weight']
+
+        assert num_frames == len(frame_weight), f'The number of frames ' \
+            f'({num_frames}) and the length of weights for each frame ' \
+            f'({len(frame_weight)}) must match'
+
+        if self.concat_tensors:
+            features = [self.backbone(torch.cat(imgs, 0))]
+        else:
+            features = [self.backbone(img) for img in imgs]
+
+        if self.with_neck:
+            features = self.neck(features, frame_weight=frame_weight)
+
+        if self.with_keypoint:
+            output = self.keypoint_head(features)
+
+        # if return loss
+        losses = dict()
+        if self.with_keypoint:
+            keypoint_losses = self.keypoint_head.get_loss(
+                output, target, target_weight)
+            losses.update(keypoint_losses)
+            keypoint_accuracy = self.keypoint_head.get_accuracy(
+                output, target, target_weight)
+            losses.update(keypoint_accuracy)
+
+        return losses
+
+    def forward_test(self, imgs, img_metas, return_heatmap=False, **kwargs):
+        """Defines the computation performed at every call when testing."""
+        # imgs (list[Fxtorch.Tensor[NxCximgHximgW]]): multiple input frames
+        assert imgs[0].size(0) == len(img_metas)
+        num_frames = len(imgs)
+        frame_weight = img_metas[0]['frame_weight']
+
+        assert num_frames == len(frame_weight), f'The number of frames ' \
+            f'({num_frames}) and the length of weights for each frame ' \
+            f'({len(frame_weight)}) must match'
+
+        batch_size, _, img_height, img_width = imgs[0].shape
+
+        if batch_size > 1:
+            assert 'bbox_id' in img_metas[0]
+
+        result = {}
+
+        if self.concat_tensors:
+            features = [self.backbone(torch.cat(imgs, 0))]
+        else:
+            features = [self.backbone(img) for img in imgs]
+
+        if self.with_neck:
+            features = self.neck(features, frame_weight=frame_weight)
+
+        if self.with_keypoint:
+            output_heatmap = self.keypoint_head.inference_model(
+                features, flip_pairs=None)
+
+        if self.test_cfg.get('flip_test', True):
+            imgs_flipped = [img.flip(3) for img in imgs]
+
+            if self.concat_tensors:
+                features_flipped = [self.backbone(torch.cat(imgs_flipped, 0))]
+            else:
+                features_flipped = [
+                    self.backbone(img_flipped) for img_flipped in imgs_flipped
+                ]
+
+            if self.with_neck:
+                features_flipped = self.neck(
+                    features_flipped, frame_weight=frame_weight)
+
+            if self.with_keypoint:
+                output_flipped_heatmap = self.keypoint_head.inference_model(
+                    features_flipped, img_metas[0]['flip_pairs'])
+                output_heatmap = (output_heatmap +
+                                  output_flipped_heatmap) * 0.5
+
+        if self.with_keypoint:
+            keypoint_result = self.keypoint_head.decode(
+                img_metas, output_heatmap, img_size=[img_width, img_height])
+            result.update(keypoint_result)
+
+            if not return_heatmap:
+                output_heatmap = None
+
+            result['output_heatmap'] = output_heatmap
+
+        return result
+
+    def forward_dummy(self, img):
+        """Used for computing network FLOPs.
+
+        See ``tools/get_flops.py``.
+
+        Args:
+            img (torch.Tensor[N,C,imgH,imgW], or list|tuple of tensors):
+                multiple input frames, N >= 2.
+
+        Returns:
+            Tensor: Output heatmaps.
+        """
+        # concat tensors if they are in a list
+        if isinstance(img, (list, tuple)):
+            img = torch.cat(img, 0)
+
+        batch_size = img.size(0)
+        assert batch_size > 1, 'Input batch size to PoseWarper ' \
+            'should be larger than 1.'
+        if batch_size == 2:
+            warnings.warn('Current batch size: 2, for pytorch2onnx and '
+                          'getting flops both.')
+        else:
+            warnings.warn(
+                f'Current batch size: {batch_size}, for getting flops only.')
+
+        frame_weight = np.random.uniform(0, 1, batch_size)
+        output = [self.backbone(img)]
+
+        if self.with_neck:
+            output = self.neck(output, frame_weight=frame_weight)
+        if self.with_keypoint:
+            output = self.keypoint_head(output)
+        return output

third_party/hamer/third-party/ViTPose/mmpose/models/detectors/top_down.py

@@ -0,0 +1,307 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import warnings
+
+import mmcv
+import numpy as np
+from mmcv.image import imwrite
+from mmcv.utils.misc import deprecated_api_warning
+from mmcv.visualization.image import imshow
+
+from mmpose.core import imshow_bboxes, imshow_keypoints
+from .. import builder
+from ..builder import POSENETS
+from .base import BasePose
+
+try:
+    from mmcv.runner import auto_fp16
+except ImportError:
+    warnings.warn('auto_fp16 from mmpose will be deprecated from v0.15.0'
+                  'Please install mmcv>=1.1.4')
+    from mmpose.core import auto_fp16
+
+
+@POSENETS.register_module()
+class TopDown(BasePose):
+    """Top-down pose detectors.
+
+    Args:
+        backbone (dict): Backbone modules to extract feature.
+        keypoint_head (dict): Keypoint head to process feature.
+        train_cfg (dict): Config for training. Default: None.
+        test_cfg (dict): Config for testing. Default: None.
+        pretrained (str): Path to the pretrained models.
+        loss_pose (None): Deprecated arguments. Please use
+            `loss_keypoint` for heads instead.
+    """
+
+    def __init__(self,
+                 backbone,
+                 neck=None,
+                 keypoint_head=None,
+                 train_cfg=None,
+                 test_cfg=None,
+                 pretrained=None,
+                 loss_pose=None):
+        super().__init__()
+        self.fp16_enabled = False
+
+        self.backbone = builder.build_backbone(backbone)
+
+        self.train_cfg = train_cfg
+        self.test_cfg = test_cfg
+
+        if neck is not None:
+            self.neck = builder.build_neck(neck)
+
+        if keypoint_head is not None:
+            keypoint_head['train_cfg'] = train_cfg
+            keypoint_head['test_cfg'] = test_cfg
+
+            if 'loss_keypoint' not in keypoint_head and loss_pose is not None:
+                warnings.warn(
+                    '`loss_pose` for TopDown is deprecated, '
+                    'use `loss_keypoint` for heads instead. See '
+                    'https://github.com/open-mmlab/mmpose/pull/382'
+                    ' for more information.', DeprecationWarning)
+                keypoint_head['loss_keypoint'] = loss_pose
+
+            self.keypoint_head = builder.build_head(keypoint_head)
+
+        self.init_weights(pretrained=pretrained)
+
+    @property
+    def with_neck(self):
+        """Check if has neck."""
+        return hasattr(self, 'neck')
+
+    @property
+    def with_keypoint(self):
+        """Check if has keypoint_head."""
+        return hasattr(self, 'keypoint_head')
+
+    def init_weights(self, pretrained=None):
+        """Weight initialization for model."""
+        self.backbone.init_weights(pretrained)
+        if self.with_neck:
+            self.neck.init_weights()
+        if self.with_keypoint:
+            self.keypoint_head.init_weights()
+
+    @auto_fp16(apply_to=('img', ))
+    def forward(self,
+                img,
+                target=None,
+                target_weight=None,
+                img_metas=None,
+                return_loss=True,
+                return_heatmap=False,
+                **kwargs):
+        """Calls either forward_train or forward_test depending on whether
+        return_loss=True. Note this setting will change the expected inputs.
+        When `return_loss=True`, img and img_meta are single-nested (i.e.
+        Tensor and List[dict]), and when `resturn_loss=False`, img and img_meta
+        should be double nested (i.e.  List[Tensor], List[List[dict]]), with
+        the outer list indicating test time augmentations.
+
+        Note:
+            - batch_size: N
+            - num_keypoints: K
+            - num_img_channel: C (Default: 3)
+            - img height: imgH
+            - img width: imgW
+            - heatmaps height: H
+            - heatmaps weight: W
+
+        Args:
+            img (torch.Tensor[NxCximgHximgW]): Input images.
+            target (torch.Tensor[NxKxHxW]): Target heatmaps.
+            target_weight (torch.Tensor[NxKx1]): Weights across
+                different joint types.
+            img_metas (list(dict)): Information about data augmentation
+                By default this includes:
+
+                - "image_file: path to the image file
+                - "center": center of the bbox
+                - "scale": scale of the bbox
+                - "rotation": rotation of the bbox
+                - "bbox_score": score of bbox
+            return_loss (bool): Option to `return loss`. `return loss=True`
+                for training, `return loss=False` for validation & test.
+            return_heatmap (bool) : Option to return heatmap.
+
+        Returns:
+            dict|tuple: if `return loss` is true, then return losses. \
+                Otherwise, return predicted poses, boxes, image paths \
+                and heatmaps.
+        """
+        if return_loss:
+            return self.forward_train(img, target, target_weight, img_metas,
+                                      **kwargs)
+        return self.forward_test(
+            img, img_metas, return_heatmap=return_heatmap, **kwargs)
+
+    def forward_train(self, img, target, target_weight, img_metas, **kwargs):
+        """Defines the computation performed at every call when training."""
+        output = self.backbone(img)
+        if self.with_neck:
+            output = self.neck(output)
+        if self.with_keypoint:
+            output = self.keypoint_head(output)
+
+        # if return loss
+        losses = dict()
+        if self.with_keypoint:
+            keypoint_losses = self.keypoint_head.get_loss(
+                output, target, target_weight)
+            losses.update(keypoint_losses)
+            keypoint_accuracy = self.keypoint_head.get_accuracy(
+                output, target, target_weight)
+            losses.update(keypoint_accuracy)
+
+        return losses
+
+    def forward_test(self, img, img_metas, return_heatmap=False, **kwargs):
+        """Defines the computation performed at every call when testing."""
+        assert img.size(0) == len(img_metas)
+        batch_size, _, img_height, img_width = img.shape
+        if batch_size > 1:
+            assert 'bbox_id' in img_metas[0]
+
+        result = {}
+
+        features = self.backbone(img)
+        if self.with_neck:
+            features = self.neck(features)
+        if self.with_keypoint:
+            output_heatmap = self.keypoint_head.inference_model(
+                features, flip_pairs=None)
+
+        if self.test_cfg.get('flip_test', True):
+            img_flipped = img.flip(3)
+            features_flipped = self.backbone(img_flipped)
+            if self.with_neck:
+                features_flipped = self.neck(features_flipped)
+            if self.with_keypoint:
+                output_flipped_heatmap = self.keypoint_head.inference_model(
+                    features_flipped, img_metas[0]['flip_pairs'])
+                output_heatmap = (output_heatmap +
+                                  output_flipped_heatmap) * 0.5
+
+        if self.with_keypoint:
+            keypoint_result = self.keypoint_head.decode(
+                img_metas, output_heatmap, img_size=[img_width, img_height])
+            result.update(keypoint_result)
+
+            if not return_heatmap:
+                output_heatmap = None
+
+            result['output_heatmap'] = output_heatmap
+
+        return result
+
+    def forward_dummy(self, img):
+        """Used for computing network FLOPs.
+
+        See ``tools/get_flops.py``.
+
+        Args:
+            img (torch.Tensor): Input image.
+
+        Returns:
+            Tensor: Output heatmaps.
+        """
+        output = self.backbone(img)
+        if self.with_neck:
+            output = self.neck(output)
+        if self.with_keypoint:
+            output = self.keypoint_head(output)
+        return output
+
+    @deprecated_api_warning({'pose_limb_color': 'pose_link_color'},
+                            cls_name='TopDown')
+    def show_result(self,
+                    img,
+                    result,
+                    skeleton=None,
+                    kpt_score_thr=0.3,
+                    bbox_color='green',
+                    pose_kpt_color=None,
+                    pose_link_color=None,
+                    text_color='white',
+                    radius=4,
+                    thickness=1,
+                    font_scale=0.5,
+                    bbox_thickness=1,
+                    win_name='',
+                    show=False,
+                    show_keypoint_weight=False,
+                    wait_time=0,
+                    out_file=None):
+        """Draw `result` over `img`.
+
+        Args:
+            img (str or Tensor): The image to be displayed.
+            result (list[dict]): The results to draw over `img`
+                (bbox_result, pose_result).
+            skeleton (list[list]): The connection of keypoints.
+                skeleton is 0-based indexing.
+            kpt_score_thr (float, optional): Minimum score of keypoints
+                to be shown. Default: 0.3.
+            bbox_color (str or tuple or :obj:`Color`): Color of bbox lines.
+            pose_kpt_color (np.array[Nx3]`): Color of N keypoints.
+                If None, do not draw keypoints.
+            pose_link_color (np.array[Mx3]): Color of M links.
+                If None, do not draw links.
+            text_color (str or tuple or :obj:`Color`): Color of texts.
+            radius (int): Radius of circles.
+            thickness (int): Thickness of lines.
+            font_scale (float): Font scales of texts.
+            win_name (str): The window name.
+            show (bool): Whether to show the image. Default: False.
+            show_keypoint_weight (bool): Whether to change the transparency
+                using the predicted confidence scores of keypoints.
+            wait_time (int): Value of waitKey param.
+                Default: 0.
+            out_file (str or None): The filename to write the image.
+                Default: None.
+
+        Returns:
+            Tensor: Visualized img, only if not `show` or `out_file`.
+        """
+        img = mmcv.imread(img)
+        img = img.copy()
+
+        bbox_result = []
+        bbox_labels = []
+        pose_result = []
+        for res in result:
+            if 'bbox' in res:
+                bbox_result.append(res['bbox'])
+                bbox_labels.append(res.get('label', None))
+            pose_result.append(res['keypoints'])
+
+        if bbox_result:
+            bboxes = np.vstack(bbox_result)
+            # draw bounding boxes
+            imshow_bboxes(
+                img,
+                bboxes,
+                labels=bbox_labels,
+                colors=bbox_color,
+                text_color=text_color,
+                thickness=bbox_thickness,
+                font_scale=font_scale,
+                show=False)
+
+        if pose_result:
+            imshow_keypoints(img, pose_result, skeleton, kpt_score_thr,
+                             pose_kpt_color, pose_link_color, radius,
+                             thickness)
+
+        if show:
+            imshow(img, win_name, wait_time)
+
+        if out_file is not None:
+            imwrite(img, out_file)
+
+        return img

third_party/hamer/third-party/ViTPose/mmpose/models/detectors/top_down_moe.py

@@ -0,0 +1,351 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import warnings
+
+import torch
+import torch.nn as nn
+
+import mmcv
+import numpy as np
+from mmcv.image import imwrite
+from mmcv.utils.misc import deprecated_api_warning
+from mmcv.visualization.image import imshow
+
+from mmpose.core import imshow_bboxes, imshow_keypoints
+from .. import builder
+from ..builder import POSENETS
+from .base import BasePose
+
+try:
+    from mmcv.runner import auto_fp16
+except ImportError:
+    warnings.warn('auto_fp16 from mmpose will be deprecated from v0.15.0'
+                  'Please install mmcv>=1.1.4')
+    from mmpose.core import auto_fp16
+
+
+@POSENETS.register_module()
+class TopDownMoE(BasePose):
+    """Top-down pose detectors.
+
+    Args:
+        backbone (dict): Backbone modules to extract feature.
+        keypoint_head (dict): Keypoint head to process feature.
+        train_cfg (dict): Config for training. Default: None.
+        test_cfg (dict): Config for testing. Default: None.
+        pretrained (str): Path to the pretrained models.
+        loss_pose (None): Deprecated arguments. Please use
+            `loss_keypoint` for heads instead.
+    """
+
+    def __init__(self,
+                 backbone,
+                 neck=None,
+                 keypoint_head=None,
+                 associate_keypoint_head=None,
+                 train_cfg=None,
+                 test_cfg=None,
+                 pretrained=None,
+                 loss_pose=None):
+        super().__init__()
+        self.fp16_enabled = False
+
+        self.backbone = builder.build_backbone(backbone)
+
+        self.train_cfg = train_cfg
+        self.test_cfg = test_cfg
+
+        if neck is not None:
+            self.neck = builder.build_neck(neck)
+
+        if keypoint_head is not None:
+            keypoint_head['train_cfg'] = train_cfg
+            keypoint_head['test_cfg'] = test_cfg
+
+            if 'loss_keypoint' not in keypoint_head and loss_pose is not None:
+                warnings.warn(
+                    '`loss_pose` for TopDown is deprecated, '
+                    'use `loss_keypoint` for heads instead. See '
+                    'https://github.com/open-mmlab/mmpose/pull/382'
+                    ' for more information.', DeprecationWarning)
+                keypoint_head['loss_keypoint'] = loss_pose
+                
+            self.keypoint_head = builder.build_head(keypoint_head)
+
+
+        associate_keypoint_heads = []
+        keypoint_heads_cnt = 1
+
+        if associate_keypoint_head is not None:
+            if not isinstance(associate_keypoint_head, list):
+                associate_keypoint_head = [associate_keypoint_head]
+            for single_keypoint_head in associate_keypoint_head:
+                single_keypoint_head['train_cfg'] = train_cfg
+                single_keypoint_head['test_cfg'] = test_cfg
+                associate_keypoint_heads.append(builder.build_head(single_keypoint_head))
+                keypoint_heads_cnt += 1
+
+        self.associate_keypoint_heads = nn.ModuleList(associate_keypoint_heads)
+
+        self.keypoint_heads_cnt = keypoint_heads_cnt
+
+        self.init_weights(pretrained=pretrained)
+
+    @property
+    def with_neck(self):
+        """Check if has neck."""
+        return hasattr(self, 'neck')
+
+    @property
+    def with_keypoint(self):
+        """Check if has keypoint_head."""
+        return hasattr(self, 'keypoint_head')
+
+    def init_weights(self, pretrained=None):
+        """Weight initialization for model."""
+        self.backbone.init_weights(pretrained)
+        if self.with_neck:
+            self.neck.init_weights()
+        if self.with_keypoint:
+            self.keypoint_head.init_weights()
+        for item in self.associate_keypoint_heads:
+            item.init_weights()
+
+    @auto_fp16(apply_to=('img', ))
+    def forward(self,
+                img,
+                target=None,
+                target_weight=None,
+                img_metas=None,
+                return_loss=True,
+                return_heatmap=False,
+                **kwargs):
+        """Calls either forward_train or forward_test depending on whether
+        return_loss=True. Note this setting will change the expected inputs.
+        When `return_loss=True`, img and img_meta are single-nested (i.e.
+        Tensor and List[dict]), and when `resturn_loss=False`, img and img_meta
+        should be double nested (i.e.  List[Tensor], List[List[dict]]), with
+        the outer list indicating test time augmentations.
+
+        Note:
+            - batch_size: N
+            - num_keypoints: K
+            - num_img_channel: C (Default: 3)
+            - img height: imgH
+            - img width: imgW
+            - heatmaps height: H
+            - heatmaps weight: W
+
+        Args:
+            img (torch.Tensor[NxCximgHximgW]): Input images.
+            target (torch.Tensor[NxKxHxW]): Target heatmaps.
+            target_weight (torch.Tensor[NxKx1]): Weights across
+                different joint types.
+            img_metas (list(dict)): Information about data augmentation
+                By default this includes:
+
+                - "image_file: path to the image file
+                - "center": center of the bbox
+                - "scale": scale of the bbox
+                - "rotation": rotation of the bbox
+                - "bbox_score": score of bbox
+            return_loss (bool): Option to `return loss`. `return loss=True`
+                for training, `return loss=False` for validation & test.
+            return_heatmap (bool) : Option to return heatmap.
+
+        Returns:
+            dict|tuple: if `return loss` is true, then return losses. \
+                Otherwise, return predicted poses, boxes, image paths \
+                and heatmaps.
+        """
+        if return_loss:
+            return self.forward_train(img, target, target_weight, img_metas,
+                                      **kwargs)
+        return self.forward_test(
+            img, img_metas, return_heatmap=return_heatmap, **kwargs)
+
+    def forward_train(self, img, target, target_weight, img_metas, **kwargs):
+        """Defines the computation performed at every call when training."""
+
+        img_sources = torch.from_numpy(np.array([ele['dataset_idx'] for ele in img_metas])).to(img.device)
+
+        output = self.backbone(img, img_sources)
+        if self.with_neck:
+            output = self.neck(output)
+        # if return loss
+        losses = dict()
+
+        main_stream_select = (img_sources == 0)
+        # if torch.sum(main_stream_select) > 0:
+        output_select = self.keypoint_head(output)
+
+        target_select = target * main_stream_select.view(-1, 1, 1, 1)
+        target_weight_select = target_weight * main_stream_select.view(-1, 1, 1)
+
+        keypoint_losses = self.keypoint_head.get_loss(
+            output_select, target_select, target_weight_select)
+        losses['main_stream_loss'] = keypoint_losses['heatmap_loss']
+        keypoint_accuracy = self.keypoint_head.get_accuracy(
+            output_select, target_select, target_weight_select)
+        losses['main_stream_acc'] = keypoint_accuracy['acc_pose']
+
+        for idx in range(1, self.keypoint_heads_cnt):
+            idx_select = (img_sources == idx)
+            target_select = target * idx_select.view(-1, 1, 1, 1)
+            target_weight_select = target_weight * idx_select.view(-1, 1, 1)
+            output_select = self.associate_keypoint_heads[idx - 1](output)
+            keypoint_losses = self.associate_keypoint_heads[idx - 1].get_loss(
+                output_select, target_select, target_weight_select)
+            losses[f'{idx}_loss'] = keypoint_losses['heatmap_loss']
+            keypoint_accuracy = self.associate_keypoint_heads[idx - 1].get_accuracy(
+                output_select, target_select, target_weight_select)
+            losses[f'{idx}_acc'] = keypoint_accuracy['acc_pose']
+
+        return losses
+
+    def forward_test(self, img, img_metas, return_heatmap=False, **kwargs):
+        """Defines the computation performed at every call when testing."""
+        assert img.size(0) == len(img_metas)
+        batch_size, _, img_height, img_width = img.shape
+        if batch_size > 1:
+            assert 'bbox_id' in img_metas[0]
+
+        result = {}
+        img_sources = torch.from_numpy(np.array([ele['dataset_idx'] for ele in img_metas])).to(img.device)
+
+        features = self.backbone(img, img_sources)
+
+        if self.with_neck:
+            features = self.neck(features)
+        if self.with_keypoint:
+            output_heatmap = self.keypoint_head.inference_model(
+                features, flip_pairs=None)
+
+        if self.test_cfg.get('flip_test', True):
+            img_flipped = img.flip(3)
+            features_flipped = self.backbone(img_flipped, img_sources)
+            if self.with_neck:
+                features_flipped = self.neck(features_flipped)
+            if self.with_keypoint:
+                output_flipped_heatmap = self.keypoint_head.inference_model(
+                    features_flipped, img_metas[0]['flip_pairs'])
+                output_heatmap = (output_heatmap +
+                                  output_flipped_heatmap) * 0.5
+
+        if self.with_keypoint:
+            keypoint_result = self.keypoint_head.decode(
+                img_metas, output_heatmap, img_size=[img_width, img_height])
+            result.update(keypoint_result)
+
+            if not return_heatmap:
+                output_heatmap = None
+
+            result['output_heatmap'] = output_heatmap
+
+        return result
+
+    def forward_dummy(self, img):
+        """Used for computing network FLOPs.
+
+        See ``tools/get_flops.py``.
+
+        Args:
+            img (torch.Tensor): Input image.
+
+        Returns:
+            Tensor: Output heatmaps.
+        """
+        output = self.backbone(img)
+        if self.with_neck:
+            output = self.neck(output)
+        if self.with_keypoint:
+            output = self.keypoint_head(output)
+        return output
+
+    @deprecated_api_warning({'pose_limb_color': 'pose_link_color'},
+                            cls_name='TopDown')
+    def show_result(self,
+                    img,
+                    result,
+                    skeleton=None,
+                    kpt_score_thr=0.3,
+                    bbox_color='green',
+                    pose_kpt_color=None,
+                    pose_link_color=None,
+                    text_color='white',
+                    radius=4,
+                    thickness=1,
+                    font_scale=0.5,
+                    bbox_thickness=1,
+                    win_name='',
+                    show=False,
+                    show_keypoint_weight=False,
+                    wait_time=0,
+                    out_file=None):
+        """Draw `result` over `img`.
+
+        Args:
+            img (str or Tensor): The image to be displayed.
+            result (list[dict]): The results to draw over `img`
+                (bbox_result, pose_result).
+            skeleton (list[list]): The connection of keypoints.
+                skeleton is 0-based indexing.
+            kpt_score_thr (float, optional): Minimum score of keypoints
+                to be shown. Default: 0.3.
+            bbox_color (str or tuple or :obj:`Color`): Color of bbox lines.
+            pose_kpt_color (np.array[Nx3]`): Color of N keypoints.
+                If None, do not draw keypoints.
+            pose_link_color (np.array[Mx3]): Color of M links.
+                If None, do not draw links.
+            text_color (str or tuple or :obj:`Color`): Color of texts.
+            radius (int): Radius of circles.
+            thickness (int): Thickness of lines.
+            font_scale (float): Font scales of texts.
+            win_name (str): The window name.
+            show (bool): Whether to show the image. Default: False.
+            show_keypoint_weight (bool): Whether to change the transparency
+                using the predicted confidence scores of keypoints.
+            wait_time (int): Value of waitKey param.
+                Default: 0.
+            out_file (str or None): The filename to write the image.
+                Default: None.
+
+        Returns:
+            Tensor: Visualized img, only if not `show` or `out_file`.
+        """
+        img = mmcv.imread(img)
+        img = img.copy()
+
+        bbox_result = []
+        bbox_labels = []
+        pose_result = []
+        for res in result:
+            if 'bbox' in res:
+                bbox_result.append(res['bbox'])
+                bbox_labels.append(res.get('label', None))
+            pose_result.append(res['keypoints'])
+
+        if bbox_result:
+            bboxes = np.vstack(bbox_result)
+            # draw bounding boxes
+            imshow_bboxes(
+                img,
+                bboxes,
+                labels=bbox_labels,
+                colors=bbox_color,
+                text_color=text_color,
+                thickness=bbox_thickness,
+                font_scale=font_scale,
+                show=False)
+
+        if pose_result:
+            imshow_keypoints(img, pose_result, skeleton, kpt_score_thr,
+                             pose_kpt_color, pose_link_color, radius,
+                             thickness)
+
+        if show:
+            imshow(img, win_name, wait_time)
+
+        if out_file is not None:
+            imwrite(img, out_file)
+
+        return img