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/datasets/pipelines/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from .bottom_up_transform import *  # noqa
+from .hand_transform import *  # noqa
+from .loading import LoadImageFromFile  # noqa
+from .mesh_transform import *  # noqa
+from .pose3d_transform import *  # noqa
+from .shared_transform import *  # noqa
+from .top_down_transform import *  # noqa

third_party/hamer/third-party/ViTPose/mmpose/datasets/pipelines/bottom_up_transform.py

@@ -0,0 +1,816 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import cv2
+import numpy as np
+
+from mmpose.core.post_processing import (get_affine_transform, get_warp_matrix,
+                                         warp_affine_joints)
+from mmpose.datasets.builder import PIPELINES
+from .shared_transform import Compose
+
+
+def _ceil_to_multiples_of(x, base=64):
+    """Transform x to the integral multiple of the base."""
+    return int(np.ceil(x / base)) * base
+
+
+def _get_multi_scale_size(image,
+                          input_size,
+                          current_scale,
+                          min_scale,
+                          use_udp=False):
+    """Get the size for multi-scale training.
+
+    Args:
+        image: Input image.
+        input_size (np.ndarray[2]): Size (w, h) of the image input.
+        current_scale (float): Scale factor.
+        min_scale (float): Minimal scale.
+        use_udp (bool): To use unbiased data processing.
+            Paper ref: Huang et al. The Devil is in the Details: Delving into
+            Unbiased Data Processing for Human Pose Estimation (CVPR 2020).
+
+    Returns:
+        tuple: A tuple containing multi-scale sizes.
+
+        - (w_resized, h_resized) (tuple(int)): resized width/height
+        - center (np.ndarray)image center
+        - scale (np.ndarray): scales wrt width/height
+    """
+    assert len(input_size) == 2
+    h, w, _ = image.shape
+
+    # calculate the size for min_scale
+    min_input_w = _ceil_to_multiples_of(min_scale * input_size[0], 64)
+    min_input_h = _ceil_to_multiples_of(min_scale * input_size[1], 64)
+    if w < h:
+        w_resized = int(min_input_w * current_scale / min_scale)
+        h_resized = int(
+            _ceil_to_multiples_of(min_input_w / w * h, 64) * current_scale /
+            min_scale)
+        if use_udp:
+            scale_w = w - 1.0
+            scale_h = (h_resized - 1.0) / (w_resized - 1.0) * (w - 1.0)
+        else:
+            scale_w = w / 200.0
+            scale_h = h_resized / w_resized * w / 200.0
+    else:
+        h_resized = int(min_input_h * current_scale / min_scale)
+        w_resized = int(
+            _ceil_to_multiples_of(min_input_h / h * w, 64) * current_scale /
+            min_scale)
+        if use_udp:
+            scale_h = h - 1.0
+            scale_w = (w_resized - 1.0) / (h_resized - 1.0) * (h - 1.0)
+        else:
+            scale_h = h / 200.0
+            scale_w = w_resized / h_resized * h / 200.0
+    if use_udp:
+        center = (scale_w / 2.0, scale_h / 2.0)
+    else:
+        center = np.array([round(w / 2.0), round(h / 2.0)])
+    return (w_resized, h_resized), center, np.array([scale_w, scale_h])
+
+
+def _resize_align_multi_scale(image, input_size, current_scale, min_scale):
+    """Resize the images for multi-scale training.
+
+    Args:
+        image: Input image
+        input_size (np.ndarray[2]): Size (w, h) of the image input
+        current_scale (float): Current scale
+        min_scale (float): Minimal scale
+
+    Returns:
+        tuple: A tuple containing image info.
+
+        - image_resized (np.ndarray): resized image
+        - center (np.ndarray): center of image
+        - scale (np.ndarray): scale
+    """
+    assert len(input_size) == 2
+    size_resized, center, scale = _get_multi_scale_size(
+        image, input_size, current_scale, min_scale)
+
+    trans = get_affine_transform(center, scale, 0, size_resized)
+    image_resized = cv2.warpAffine(image, trans, size_resized)
+
+    return image_resized, center, scale
+
+
+def _resize_align_multi_scale_udp(image, input_size, current_scale, min_scale):
+    """Resize the images for multi-scale training.
+
+    Args:
+        image: Input image
+        input_size (np.ndarray[2]): Size (w, h) of the image input
+        current_scale (float): Current scale
+        min_scale (float): Minimal scale
+
+    Returns:
+        tuple: A tuple containing image info.
+
+        - image_resized (np.ndarray): resized image
+        - center (np.ndarray): center of image
+        - scale (np.ndarray): scale
+    """
+    assert len(input_size) == 2
+    size_resized, _, _ = _get_multi_scale_size(image, input_size,
+                                               current_scale, min_scale, True)
+
+    _, center, scale = _get_multi_scale_size(image, input_size, min_scale,
+                                             min_scale, True)
+
+    trans = get_warp_matrix(
+        theta=0,
+        size_input=np.array(scale, dtype=np.float32),
+        size_dst=np.array(size_resized, dtype=np.float32) - 1.0,
+        size_target=np.array(scale, dtype=np.float32))
+    image_resized = cv2.warpAffine(
+        image.copy(), trans, size_resized, flags=cv2.INTER_LINEAR)
+
+    return image_resized, center, scale
+
+
+class HeatmapGenerator:
+    """Generate heatmaps for bottom-up models.
+
+    Args:
+        num_joints (int): Number of keypoints
+        output_size (np.ndarray): Size (w, h) of feature map
+        sigma (int): Sigma of the heatmaps.
+        use_udp (bool): To use unbiased data processing.
+            Paper ref: Huang et al. The Devil is in the Details: Delving into
+            Unbiased Data Processing for Human Pose Estimation (CVPR 2020).
+    """
+
+    def __init__(self, output_size, num_joints, sigma=-1, use_udp=False):
+        if not isinstance(output_size, np.ndarray):
+            output_size = np.array(output_size)
+        if output_size.size > 1:
+            assert len(output_size) == 2
+            self.output_size = output_size
+        else:
+            self.output_size = np.array([output_size, output_size],
+                                        dtype=np.int)
+        self.num_joints = num_joints
+        if sigma < 0:
+            sigma = self.output_size.prod()**0.5 / 64
+        self.sigma = sigma
+        size = 6 * sigma + 3
+        self.use_udp = use_udp
+        if use_udp:
+            self.x = np.arange(0, size, 1, np.float32)
+            self.y = self.x[:, None]
+        else:
+            x = np.arange(0, size, 1, np.float32)
+            y = x[:, None]
+            x0, y0 = 3 * sigma + 1, 3 * sigma + 1
+            self.g = np.exp(-((x - x0)**2 + (y - y0)**2) / (2 * sigma**2))
+
+    def __call__(self, joints):
+        """Generate heatmaps."""
+        hms = np.zeros(
+            (self.num_joints, self.output_size[1], self.output_size[0]),
+            dtype=np.float32)
+
+        sigma = self.sigma
+        for p in joints:
+            for idx, pt in enumerate(p):
+                if pt[2] > 0:
+                    x, y = int(pt[0]), int(pt[1])
+                    if x < 0 or y < 0 or \
+                       x >= self.output_size[0] or y >= self.output_size[1]:
+                        continue
+
+                    if self.use_udp:
+                        x0 = 3 * sigma + 1 + pt[0] - x
+                        y0 = 3 * sigma + 1 + pt[1] - y
+                        g = np.exp(-((self.x - x0)**2 + (self.y - y0)**2) /
+                                   (2 * sigma**2))
+                    else:
+                        g = self.g
+
+                    ul = int(np.round(x - 3 * sigma -
+                                      1)), int(np.round(y - 3 * sigma - 1))
+                    br = int(np.round(x + 3 * sigma +
+                                      2)), int(np.round(y + 3 * sigma + 2))
+
+                    c, d = max(0,
+                               -ul[0]), min(br[0], self.output_size[0]) - ul[0]
+                    a, b = max(0,
+                               -ul[1]), min(br[1], self.output_size[1]) - ul[1]
+
+                    cc, dd = max(0, ul[0]), min(br[0], self.output_size[0])
+                    aa, bb = max(0, ul[1]), min(br[1], self.output_size[1])
+                    hms[idx, aa:bb,
+                        cc:dd] = np.maximum(hms[idx, aa:bb, cc:dd], g[a:b,
+                                                                      c:d])
+        return hms
+
+
+class JointsEncoder:
+    """Encodes the visible joints into (coordinates, score); The coordinate of
+    one joint and its score are of `int` type.
+
+    (idx * output_size**2 + y * output_size + x, 1) or (0, 0).
+
+    Args:
+        max_num_people(int): Max number of people in an image
+        num_joints(int): Number of keypoints
+        output_size(np.ndarray): Size (w, h) of feature map
+        tag_per_joint(bool):  Option to use one tag map per joint.
+    """
+
+    def __init__(self, max_num_people, num_joints, output_size, tag_per_joint):
+        self.max_num_people = max_num_people
+        self.num_joints = num_joints
+        if not isinstance(output_size, np.ndarray):
+            output_size = np.array(output_size)
+        if output_size.size > 1:
+            assert len(output_size) == 2
+            self.output_size = output_size
+        else:
+            self.output_size = np.array([output_size, output_size],
+                                        dtype=np.int)
+        self.tag_per_joint = tag_per_joint
+
+    def __call__(self, joints):
+        """
+        Note:
+            - number of people in image: N
+            - number of keypoints: K
+            - max number of people in an image: M
+
+        Args:
+            joints (np.ndarray[N,K,3])
+
+        Returns:
+            visible_kpts (np.ndarray[M,K,2]).
+        """
+        visible_kpts = np.zeros((self.max_num_people, self.num_joints, 2),
+                                dtype=np.float32)
+        for i in range(len(joints)):
+            tot = 0
+            for idx, pt in enumerate(joints[i]):
+                x, y = int(pt[0]), int(pt[1])
+                if (pt[2] > 0 and 0 <= y < self.output_size[1]
+                        and 0 <= x < self.output_size[0]):
+                    if self.tag_per_joint:
+                        visible_kpts[i][tot] = \
+                            (idx * self.output_size.prod()
+                             + y * self.output_size[0] + x, 1)
+                    else:
+                        visible_kpts[i][tot] = (y * self.output_size[0] + x, 1)
+                    tot += 1
+        return visible_kpts
+
+
+class PAFGenerator:
+    """Generate part affinity fields.
+
+    Args:
+        output_size (np.ndarray): Size (w, h) of feature map.
+        limb_width (int): Limb width of part affinity fields.
+        skeleton (list[list]): connections of joints.
+    """
+
+    def __init__(self, output_size, limb_width, skeleton):
+        if not isinstance(output_size, np.ndarray):
+            output_size = np.array(output_size)
+        if output_size.size > 1:
+            assert len(output_size) == 2
+            self.output_size = output_size
+        else:
+            self.output_size = np.array([output_size, output_size],
+                                        dtype=np.int)
+        self.limb_width = limb_width
+        self.skeleton = skeleton
+
+    def _accumulate_paf_map_(self, pafs, src, dst, count):
+        """Accumulate part affinity fields between two given joints.
+
+        Args:
+            pafs (np.ndarray[2,H,W]): paf maps (2 dimensions:x axis and
+                y axis) for a certain limb connection. This argument will
+                be modified inplace.
+            src (np.ndarray[2,]): coordinates of the source joint.
+            dst (np.ndarray[2,]): coordinates of the destination joint.
+            count (np.ndarray[H,W]): count map that preserves the number
+                of non-zero vectors at each point. This argument will be
+                modified inplace.
+        """
+        limb_vec = dst - src
+        norm = np.linalg.norm(limb_vec)
+        if norm == 0:
+            unit_limb_vec = np.zeros(2)
+        else:
+            unit_limb_vec = limb_vec / norm
+
+        min_x = max(np.floor(min(src[0], dst[0]) - self.limb_width), 0)
+        max_x = min(
+            np.ceil(max(src[0], dst[0]) + self.limb_width),
+            self.output_size[0] - 1)
+        min_y = max(np.floor(min(src[1], dst[1]) - self.limb_width), 0)
+        max_y = min(
+            np.ceil(max(src[1], dst[1]) + self.limb_width),
+            self.output_size[1] - 1)
+
+        range_x = list(range(int(min_x), int(max_x + 1), 1))
+        range_y = list(range(int(min_y), int(max_y + 1), 1))
+
+        mask = np.zeros_like(count, dtype=bool)
+        if len(range_x) > 0 and len(range_y) > 0:
+            xx, yy = np.meshgrid(range_x, range_y)
+            delta_x = xx - src[0]
+            delta_y = yy - src[1]
+            dist = np.abs(delta_x * unit_limb_vec[1] -
+                          delta_y * unit_limb_vec[0])
+            mask_local = (dist < self.limb_width)
+            mask[yy, xx] = mask_local
+
+        pafs[0, mask] += unit_limb_vec[0]
+        pafs[1, mask] += unit_limb_vec[1]
+        count += mask
+
+        return pafs, count
+
+    def __call__(self, joints):
+        """Generate the target part affinity fields."""
+        pafs = np.zeros(
+            (len(self.skeleton) * 2, self.output_size[1], self.output_size[0]),
+            dtype=np.float32)
+
+        for idx, sk in enumerate(self.skeleton):
+            count = np.zeros((self.output_size[1], self.output_size[0]),
+                             dtype=np.float32)
+
+            for p in joints:
+                src = p[sk[0]]
+                dst = p[sk[1]]
+                if src[2] > 0 and dst[2] > 0:
+                    self._accumulate_paf_map_(pafs[2 * idx:2 * idx + 2],
+                                              src[:2], dst[:2], count)
+
+            pafs[2 * idx:2 * idx + 2] /= np.maximum(count, 1)
+
+        return pafs
+
+
+@PIPELINES.register_module()
+class BottomUpRandomFlip:
+    """Data augmentation with random image flip for bottom-up.
+
+    Args:
+        flip_prob (float): Probability of flip.
+    """
+
+    def __init__(self, flip_prob=0.5):
+        self.flip_prob = flip_prob
+
+    def __call__(self, results):
+        """Perform data augmentation with random image flip."""
+        image, mask, joints = results['img'], results['mask'], results[
+            'joints']
+        self.flip_index = results['ann_info']['flip_index']
+        self.output_size = results['ann_info']['heatmap_size']
+
+        assert isinstance(mask, list)
+        assert isinstance(joints, list)
+        assert len(mask) == len(joints)
+        assert len(mask) == len(self.output_size)
+
+        if np.random.random() < self.flip_prob:
+            image = image[:, ::-1].copy() - np.zeros_like(image)
+            for i, _output_size in enumerate(self.output_size):
+                if not isinstance(_output_size, np.ndarray):
+                    _output_size = np.array(_output_size)
+                if _output_size.size > 1:
+                    assert len(_output_size) == 2
+                else:
+                    _output_size = np.array([_output_size, _output_size],
+                                            dtype=np.int)
+                mask[i] = mask[i][:, ::-1].copy()
+                joints[i] = joints[i][:, self.flip_index]
+                joints[i][:, :, 0] = _output_size[0] - joints[i][:, :, 0] - 1
+        results['img'], results['mask'], results[
+            'joints'] = image, mask, joints
+        return results
+
+
+@PIPELINES.register_module()
+class BottomUpRandomAffine:
+    """Data augmentation with random scaling & rotating.
+
+    Args:
+        rot_factor (int): Rotating to [-rotation_factor, rotation_factor]
+        scale_factor (float): Scaling to [1-scale_factor, 1+scale_factor]
+        scale_type: wrt ``long`` or ``short`` length of the image.
+        trans_factor: Translation factor.
+        use_udp (bool): To use unbiased data processing.
+            Paper ref: Huang et al. The Devil is in the Details: Delving into
+            Unbiased Data Processing for Human Pose Estimation (CVPR 2020).
+    """
+
+    def __init__(self,
+                 rot_factor,
+                 scale_factor,
+                 scale_type,
+                 trans_factor,
+                 use_udp=False):
+        self.max_rotation = rot_factor
+        self.min_scale = scale_factor[0]
+        self.max_scale = scale_factor[1]
+        self.scale_type = scale_type
+        self.trans_factor = trans_factor
+        self.use_udp = use_udp
+
+    def _get_scale(self, image_size, resized_size):
+        w, h = image_size
+        w_resized, h_resized = resized_size
+        if w / w_resized < h / h_resized:
+            if self.scale_type == 'long':
+                w_pad = h / h_resized * w_resized
+                h_pad = h
+            elif self.scale_type == 'short':
+                w_pad = w
+                h_pad = w / w_resized * h_resized
+            else:
+                raise ValueError(f'Unknown scale type: {self.scale_type}')
+        else:
+            if self.scale_type == 'long':
+                w_pad = w
+                h_pad = w / w_resized * h_resized
+            elif self.scale_type == 'short':
+                w_pad = h / h_resized * w_resized
+                h_pad = h
+            else:
+                raise ValueError(f'Unknown scale type: {self.scale_type}')
+
+        scale = np.array([w_pad, h_pad], dtype=np.float32)
+
+        return scale
+
+    def __call__(self, results):
+        """Perform data augmentation with random scaling & rotating."""
+        image, mask, joints = results['img'], results['mask'], results[
+            'joints']
+
+        self.input_size = results['ann_info']['image_size']
+        if not isinstance(self.input_size, np.ndarray):
+            self.input_size = np.array(self.input_size)
+        if self.input_size.size > 1:
+            assert len(self.input_size) == 2
+        else:
+            self.input_size = [self.input_size, self.input_size]
+        self.output_size = results['ann_info']['heatmap_size']
+
+        assert isinstance(mask, list)
+        assert isinstance(joints, list)
+        assert len(mask) == len(joints)
+        assert len(mask) == len(self.output_size), (len(mask),
+                                                    len(self.output_size),
+                                                    self.output_size)
+
+        height, width = image.shape[:2]
+        if self.use_udp:
+            center = np.array(((width - 1.0) / 2, (height - 1.0) / 2))
+        else:
+            center = np.array((width / 2, height / 2))
+
+        img_scale = np.array([width, height], dtype=np.float32)
+        aug_scale = np.random.random() * (self.max_scale - self.min_scale) \
+            + self.min_scale
+        img_scale *= aug_scale
+        aug_rot = (np.random.random() * 2 - 1) * self.max_rotation
+
+        if self.trans_factor > 0:
+            dx = np.random.randint(-self.trans_factor * img_scale[0] / 200.0,
+                                   self.trans_factor * img_scale[0] / 200.0)
+            dy = np.random.randint(-self.trans_factor * img_scale[1] / 200.0,
+                                   self.trans_factor * img_scale[1] / 200.0)
+
+            center[0] += dx
+            center[1] += dy
+        if self.use_udp:
+            for i, _output_size in enumerate(self.output_size):
+                if not isinstance(_output_size, np.ndarray):
+                    _output_size = np.array(_output_size)
+                if _output_size.size > 1:
+                    assert len(_output_size) == 2
+                else:
+                    _output_size = [_output_size, _output_size]
+
+                scale = self._get_scale(img_scale, _output_size)
+
+                trans = get_warp_matrix(
+                    theta=aug_rot,
+                    size_input=center * 2.0,
+                    size_dst=np.array(
+                        (_output_size[0], _output_size[1]), dtype=np.float32) -
+                    1.0,
+                    size_target=scale)
+                mask[i] = cv2.warpAffine(
+                    (mask[i] * 255).astype(np.uint8),
+                    trans, (int(_output_size[0]), int(_output_size[1])),
+                    flags=cv2.INTER_LINEAR) / 255
+                mask[i] = (mask[i] > 0.5).astype(np.float32)
+                joints[i][:, :, 0:2] = \
+                    warp_affine_joints(joints[i][:, :, 0:2].copy(), trans)
+                if results['ann_info']['scale_aware_sigma']:
+                    joints[i][:, :, 3] = joints[i][:, :, 3] / aug_scale
+            scale = self._get_scale(img_scale, self.input_size)
+            mat_input = get_warp_matrix(
+                theta=aug_rot,
+                size_input=center * 2.0,
+                size_dst=np.array((self.input_size[0], self.input_size[1]),
+                                  dtype=np.float32) - 1.0,
+                size_target=scale)
+            image = cv2.warpAffine(
+                image,
+                mat_input, (int(self.input_size[0]), int(self.input_size[1])),
+                flags=cv2.INTER_LINEAR)
+        else:
+            for i, _output_size in enumerate(self.output_size):
+                if not isinstance(_output_size, np.ndarray):
+                    _output_size = np.array(_output_size)
+                if _output_size.size > 1:
+                    assert len(_output_size) == 2
+                else:
+                    _output_size = [_output_size, _output_size]
+                scale = self._get_scale(img_scale, _output_size)
+                mat_output = get_affine_transform(
+                    center=center,
+                    scale=scale / 200.0,
+                    rot=aug_rot,
+                    output_size=_output_size)
+                mask[i] = cv2.warpAffine(
+                    (mask[i] * 255).astype(np.uint8), mat_output,
+                    (int(_output_size[0]), int(_output_size[1]))) / 255
+                mask[i] = (mask[i] > 0.5).astype(np.float32)
+
+                joints[i][:, :, 0:2] = \
+                    warp_affine_joints(joints[i][:, :, 0:2], mat_output)
+                if results['ann_info']['scale_aware_sigma']:
+                    joints[i][:, :, 3] = joints[i][:, :, 3] / aug_scale
+
+            scale = self._get_scale(img_scale, self.input_size)
+            mat_input = get_affine_transform(
+                center=center,
+                scale=scale / 200.0,
+                rot=aug_rot,
+                output_size=self.input_size)
+            image = cv2.warpAffine(image, mat_input, (int(
+                self.input_size[0]), int(self.input_size[1])))
+
+        results['img'], results['mask'], results[
+            'joints'] = image, mask, joints
+
+        return results
+
+
+@PIPELINES.register_module()
+class BottomUpGenerateHeatmapTarget:
+    """Generate multi-scale heatmap target for bottom-up.
+
+    Args:
+        sigma (int): Sigma of heatmap Gaussian
+        max_num_people (int): Maximum number of people in an image
+        use_udp (bool): To use unbiased data processing.
+            Paper ref: Huang et al. The Devil is in the Details: Delving into
+            Unbiased Data Processing for Human Pose Estimation (CVPR 2020).
+    """
+
+    def __init__(self, sigma, use_udp=False):
+        self.sigma = sigma
+        self.use_udp = use_udp
+
+    def _generate(self, num_joints, heatmap_size):
+        """Get heatmap generator."""
+        heatmap_generator = [
+            HeatmapGenerator(output_size, num_joints, self.sigma, self.use_udp)
+            for output_size in heatmap_size
+        ]
+        return heatmap_generator
+
+    def __call__(self, results):
+        """Generate multi-scale heatmap target for bottom-up."""
+        heatmap_generator = \
+            self._generate(results['ann_info']['num_joints'],
+                           results['ann_info']['heatmap_size'])
+        target_list = list()
+        joints_list = results['joints']
+
+        for scale_id in range(results['ann_info']['num_scales']):
+            heatmaps = heatmap_generator[scale_id](joints_list[scale_id])
+            target_list.append(heatmaps.astype(np.float32))
+        results['target'] = target_list
+
+        return results
+
+
+@PIPELINES.register_module()
+class BottomUpGenerateTarget:
+    """Generate multi-scale heatmap target for associate embedding.
+
+    Args:
+        sigma (int): Sigma of heatmap Gaussian
+        max_num_people (int): Maximum number of people in an image
+        use_udp (bool): To use unbiased data processing.
+            Paper ref: Huang et al. The Devil is in the Details: Delving into
+            Unbiased Data Processing for Human Pose Estimation (CVPR 2020).
+    """
+
+    def __init__(self, sigma, max_num_people, use_udp=False):
+        self.sigma = sigma
+        self.max_num_people = max_num_people
+        self.use_udp = use_udp
+
+    def _generate(self, num_joints, heatmap_size):
+        """Get heatmap generator and joint encoder."""
+        heatmap_generator = [
+            HeatmapGenerator(output_size, num_joints, self.sigma, self.use_udp)
+            for output_size in heatmap_size
+        ]
+        joints_encoder = [
+            JointsEncoder(self.max_num_people, num_joints, output_size, True)
+            for output_size in heatmap_size
+        ]
+        return heatmap_generator, joints_encoder
+
+    def __call__(self, results):
+        """Generate multi-scale heatmap target for bottom-up."""
+        heatmap_generator, joints_encoder = \
+            self._generate(results['ann_info']['num_joints'],
+                           results['ann_info']['heatmap_size'])
+        target_list = list()
+        mask_list, joints_list = results['mask'], results['joints']
+
+        for scale_id in range(results['ann_info']['num_scales']):
+            target_t = heatmap_generator[scale_id](joints_list[scale_id])
+            joints_t = joints_encoder[scale_id](joints_list[scale_id])
+
+            target_list.append(target_t.astype(np.float32))
+            mask_list[scale_id] = mask_list[scale_id].astype(np.float32)
+            joints_list[scale_id] = joints_t.astype(np.int32)
+
+        results['masks'], results['joints'] = mask_list, joints_list
+        results['targets'] = target_list
+
+        return results
+
+
+@PIPELINES.register_module()
+class BottomUpGeneratePAFTarget:
+    """Generate multi-scale heatmaps and part affinity fields (PAF) target for
+    bottom-up. Paper ref: Cao et al. Realtime Multi-Person 2D Human Pose
+    Estimation using Part Affinity Fields (CVPR 2017).
+
+    Args:
+        limb_width (int): Limb width of part affinity fields
+    """
+
+    def __init__(self, limb_width, skeleton=None):
+        self.limb_width = limb_width
+        self.skeleton = skeleton
+
+    def _generate(self, heatmap_size, skeleton):
+        """Get PAF generator."""
+        paf_generator = [
+            PAFGenerator(output_size, self.limb_width, skeleton)
+            for output_size in heatmap_size
+        ]
+        return paf_generator
+
+    def __call__(self, results):
+        """Generate multi-scale part affinity fields for bottom-up."""
+        if self.skeleton is None:
+            assert results['ann_info']['skeleton'] is not None
+            self.skeleton = results['ann_info']['skeleton']
+
+        paf_generator = \
+            self._generate(results['ann_info']['heatmap_size'],
+                           self.skeleton)
+        target_list = list()
+        joints_list = results['joints']
+
+        for scale_id in range(results['ann_info']['num_scales']):
+            pafs = paf_generator[scale_id](joints_list[scale_id])
+            target_list.append(pafs.astype(np.float32))
+
+        results['target'] = target_list
+
+        return results
+
+
+@PIPELINES.register_module()
+class BottomUpGetImgSize:
+    """Get multi-scale image sizes for bottom-up, including base_size and
+    test_scale_factor. Keep the ratio and the image is resized to
+    `results['ann_info']['image_size']×current_scale`.
+
+    Args:
+        test_scale_factor (List[float]): Multi scale
+        current_scale (int): default 1
+        use_udp (bool): To use unbiased data processing.
+            Paper ref: Huang et al. The Devil is in the Details: Delving into
+            Unbiased Data Processing for Human Pose Estimation (CVPR 2020).
+    """
+
+    def __init__(self, test_scale_factor, current_scale=1, use_udp=False):
+        self.test_scale_factor = test_scale_factor
+        self.min_scale = min(test_scale_factor)
+        self.current_scale = current_scale
+        self.use_udp = use_udp
+
+    def __call__(self, results):
+        """Get multi-scale image sizes for bottom-up."""
+        input_size = results['ann_info']['image_size']
+        if not isinstance(input_size, np.ndarray):
+            input_size = np.array(input_size)
+        if input_size.size > 1:
+            assert len(input_size) == 2
+        else:
+            input_size = np.array([input_size, input_size], dtype=np.int)
+        img = results['img']
+
+        h, w, _ = img.shape
+
+        # calculate the size for min_scale
+        min_input_w = _ceil_to_multiples_of(self.min_scale * input_size[0], 64)
+        min_input_h = _ceil_to_multiples_of(self.min_scale * input_size[1], 64)
+        if w < h:
+            w_resized = int(min_input_w * self.current_scale / self.min_scale)
+            h_resized = int(
+                _ceil_to_multiples_of(min_input_w / w * h, 64) *
+                self.current_scale / self.min_scale)
+            if self.use_udp:
+                scale_w = w - 1.0
+                scale_h = (h_resized - 1.0) / (w_resized - 1.0) * (w - 1.0)
+            else:
+                scale_w = w / 200.0
+                scale_h = h_resized / w_resized * w / 200.0
+        else:
+            h_resized = int(min_input_h * self.current_scale / self.min_scale)
+            w_resized = int(
+                _ceil_to_multiples_of(min_input_h / h * w, 64) *
+                self.current_scale / self.min_scale)
+            if self.use_udp:
+                scale_h = h - 1.0
+                scale_w = (w_resized - 1.0) / (h_resized - 1.0) * (h - 1.0)
+            else:
+                scale_h = h / 200.0
+                scale_w = w_resized / h_resized * h / 200.0
+        if self.use_udp:
+            center = (scale_w / 2.0, scale_h / 2.0)
+        else:
+            center = np.array([round(w / 2.0), round(h / 2.0)])
+        results['ann_info']['test_scale_factor'] = self.test_scale_factor
+        results['ann_info']['base_size'] = (w_resized, h_resized)
+        results['ann_info']['center'] = center
+        results['ann_info']['scale'] = np.array([scale_w, scale_h])
+
+        return results
+
+
+@PIPELINES.register_module()
+class BottomUpResizeAlign:
+    """Resize multi-scale size and align transform for bottom-up.
+
+    Args:
+        transforms (List): ToTensor & Normalize
+        use_udp (bool): To use unbiased data processing.
+            Paper ref: Huang et al. The Devil is in the Details: Delving into
+            Unbiased Data Processing for Human Pose Estimation (CVPR 2020).
+    """
+
+    def __init__(self, transforms, use_udp=False):
+        self.transforms = Compose(transforms)
+        if use_udp:
+            self._resize_align_multi_scale = _resize_align_multi_scale_udp
+        else:
+            self._resize_align_multi_scale = _resize_align_multi_scale
+
+    def __call__(self, results):
+        """Resize multi-scale size and align transform for bottom-up."""
+        input_size = results['ann_info']['image_size']
+        if not isinstance(input_size, np.ndarray):
+            input_size = np.array(input_size)
+        if input_size.size > 1:
+            assert len(input_size) == 2
+        else:
+            input_size = np.array([input_size, input_size], dtype=np.int)
+        test_scale_factor = results['ann_info']['test_scale_factor']
+        aug_data = []
+
+        for _, s in enumerate(sorted(test_scale_factor, reverse=True)):
+            _results = results.copy()
+            image_resized, _, _ = self._resize_align_multi_scale(
+                _results['img'], input_size, s, min(test_scale_factor))
+            _results['img'] = image_resized
+            _results = self.transforms(_results)
+            transformed_img = _results['img'].unsqueeze(0)
+            aug_data.append(transformed_img)
+
+        results['ann_info']['aug_data'] = aug_data
+
+        return results

third_party/hamer/third-party/ViTPose/mmpose/datasets/pipelines/hand_transform.py

@@ -0,0 +1,63 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import numpy as np
+
+from mmpose.datasets.builder import PIPELINES
+from .top_down_transform import TopDownRandomFlip
+
+
+@PIPELINES.register_module()
+class HandRandomFlip(TopDownRandomFlip):
+    """Data augmentation with random image flip. A child class of
+    TopDownRandomFlip.
+
+    Required keys: 'img', 'joints_3d', 'joints_3d_visible', 'center',
+    'hand_type', 'rel_root_depth' and 'ann_info'.
+
+    Modifies key: 'img', 'joints_3d', 'joints_3d_visible', 'center',
+    'hand_type', 'rel_root_depth'.
+
+    Args:
+        flip_prob (float): Probability of flip.
+    """
+
+    def __call__(self, results):
+        """Perform data augmentation with random image flip."""
+        # base flip augmentation
+        super().__call__(results)
+
+        # flip hand type and root depth
+        hand_type = results['hand_type']
+        rel_root_depth = results['rel_root_depth']
+        flipped = results['flipped']
+        if flipped:
+            hand_type[0], hand_type[1] = hand_type[1], hand_type[0]
+            rel_root_depth = -rel_root_depth
+        results['hand_type'] = hand_type
+        results['rel_root_depth'] = rel_root_depth
+        return results
+
+
+@PIPELINES.register_module()
+class HandGenerateRelDepthTarget:
+    """Generate the target relative root depth.
+
+    Required keys: 'rel_root_depth', 'rel_root_valid', 'ann_info'.
+
+    Modified keys: 'target', 'target_weight'.
+    """
+
+    def __init__(self):
+        pass
+
+    def __call__(self, results):
+        """Generate the target heatmap."""
+        rel_root_depth = results['rel_root_depth']
+        rel_root_valid = results['rel_root_valid']
+        cfg = results['ann_info']
+        D = cfg['heatmap_size_root']
+        root_depth_bound = cfg['root_depth_bound']
+        target = (rel_root_depth / root_depth_bound + 0.5) * D
+        target_weight = rel_root_valid * (target >= 0) * (target <= D)
+        results['target'] = target * np.ones(1, dtype=np.float32)
+        results['target_weight'] = target_weight * np.ones(1, dtype=np.float32)
+        return results

third_party/hamer/third-party/ViTPose/mmpose/datasets/pipelines/loading.py

@@ -0,0 +1,91 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import mmcv
+import numpy as np
+
+from ..builder import PIPELINES
+
+
+@PIPELINES.register_module()
+class LoadImageFromFile:
+    """Loading image(s) from file.
+
+    Required key: "image_file".
+
+    Added key: "img".
+
+    Args:
+        to_float32 (bool): Whether to convert the loaded image to a float32
+            numpy array. If set to False, the loaded image is an uint8 array.
+            Defaults to False.
+        color_type (str): Flags specifying the color type of a loaded image,
+          candidates are 'color', 'grayscale' and 'unchanged'.
+        channel_order (str): Order of channel, candidates are 'bgr' and 'rgb'.
+        file_client_args (dict): Arguments to instantiate a FileClient.
+            See :class:`mmcv.fileio.FileClient` for details.
+            Defaults to ``dict(backend='disk')``.
+    """
+
+    def __init__(self,
+                 to_float32=False,
+                 color_type='color',
+                 channel_order='rgb',
+                 file_client_args=dict(backend='disk')):
+        self.to_float32 = to_float32
+        self.color_type = color_type
+        self.channel_order = channel_order
+        self.file_client_args = file_client_args.copy()
+        self.file_client = None
+
+    def _read_image(self, path):
+        img_bytes = self.file_client.get(path)
+        img = mmcv.imfrombytes(
+            img_bytes, flag=self.color_type, channel_order=self.channel_order)
+        if img is None:
+            raise ValueError(f'Fail to read {path}')
+        if self.to_float32:
+            img = img.astype(np.float32)
+        return img
+
+    def __call__(self, results):
+        """Loading image(s) from file."""
+        if self.file_client is None:
+            self.file_client = mmcv.FileClient(**self.file_client_args)
+
+        image_file = results.get('image_file', None)
+
+        if isinstance(image_file, (list, tuple)):
+            # Load images from a list of paths
+            results['img'] = [self._read_image(path) for path in image_file]
+        elif image_file is not None:
+            # Load single image from path
+            results['img'] = self._read_image(image_file)
+        else:
+            if 'img' not in results:
+                # If `image_file`` is not in results, check the `img` exists
+                # and format the image. This for compatibility when the image
+                # is manually set outside the pipeline.
+                raise KeyError('Either `image_file` or `img` should exist in '
+                               'results.')
+            assert isinstance(results['img'], np.ndarray)
+            if self.color_type == 'color' and self.channel_order == 'rgb':
+                # The original results['img'] is assumed to be image(s) in BGR
+                # order, so we convert the color according to the arguments.
+                if results['img'].ndim == 3:
+                    results['img'] = mmcv.bgr2rgb(results['img'])
+                elif results['img'].ndim == 4:
+                    results['img'] = np.concatenate(
+                        [mmcv.bgr2rgb(img) for img in results['img']], axis=0)
+                else:
+                    raise ValueError('results["img"] has invalid shape '
+                                     f'{results["img"].shape}')
+
+            results['image_file'] = None
+
+        return results
+
+    def __repr__(self):
+        repr_str = (f'{self.__class__.__name__}('
+                    f'to_float32={self.to_float32}, '
+                    f"color_type='{self.color_type}', "
+                    f'file_client_args={self.file_client_args})')
+        return repr_str

third_party/hamer/third-party/ViTPose/mmpose/datasets/pipelines/mesh_transform.py

@@ -0,0 +1,399 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import cv2
+import mmcv
+import numpy as np
+import torch
+
+from mmpose.core.post_processing import (affine_transform, fliplr_joints,
+                                         get_affine_transform)
+from mmpose.datasets.builder import PIPELINES
+
+
+def _flip_smpl_pose(pose):
+    """Flip SMPL pose parameters horizontally.
+
+    Args:
+        pose (np.ndarray([72])): SMPL pose parameters
+
+    Returns:
+        pose_flipped
+    """
+
+    flippedParts = [
+        0, 1, 2, 6, 7, 8, 3, 4, 5, 9, 10, 11, 15, 16, 17, 12, 13, 14, 18, 19,
+        20, 24, 25, 26, 21, 22, 23, 27, 28, 29, 33, 34, 35, 30, 31, 32, 36, 37,
+        38, 42, 43, 44, 39, 40, 41, 45, 46, 47, 51, 52, 53, 48, 49, 50, 57, 58,
+        59, 54, 55, 56, 63, 64, 65, 60, 61, 62, 69, 70, 71, 66, 67, 68
+    ]
+    pose_flipped = pose[flippedParts]
+    # Negate the second and the third dimension of the axis-angle
+    pose_flipped[1::3] = -pose_flipped[1::3]
+    pose_flipped[2::3] = -pose_flipped[2::3]
+    return pose_flipped
+
+
+def _flip_iuv(iuv, uv_type='BF'):
+    """Flip IUV image horizontally.
+
+    Note:
+        IUV image height: H
+        IUV image width: W
+
+    Args:
+        iuv np.ndarray([H, W, 3]): IUV image
+        uv_type (str): The type of the UV map.
+            Candidate values:
+                'DP': The UV map used in DensePose project.
+                'SMPL': The default UV map of SMPL model.
+                'BF': The UV map used in DecoMR project.
+            Default: 'BF'
+
+    Returns:
+        iuv_flipped np.ndarray([H, W, 3]): Flipped IUV image
+    """
+    assert uv_type in ['DP', 'SMPL', 'BF']
+    if uv_type == 'BF':
+        iuv_flipped = iuv[:, ::-1, :]
+        iuv_flipped[:, :, 1] = 255 - iuv_flipped[:, :, 1]
+    else:
+        # The flip of other UV map is complex, not finished yet.
+        raise NotImplementedError(
+            f'The flip of {uv_type} UV map is not implemented yet.')
+
+    return iuv_flipped
+
+
+def _construct_rotation_matrix(rot, size=3):
+    """Construct the in-plane rotation matrix.
+
+    Args:
+        rot (float): Rotation angle (degree).
+        size (int): The size of the rotation matrix.
+            Candidate Values: 2, 3. Defaults to 3.
+
+    Returns:
+        rot_mat (np.ndarray([size, size]): Rotation matrix.
+    """
+    rot_mat = np.eye(size, dtype=np.float32)
+    if rot != 0:
+        rot_rad = np.deg2rad(rot)
+        sn, cs = np.sin(rot_rad), np.cos(rot_rad)
+        rot_mat[0, :2] = [cs, -sn]
+        rot_mat[1, :2] = [sn, cs]
+
+    return rot_mat
+
+
+def _rotate_joints_3d(joints_3d, rot):
+    """Rotate the 3D joints in the local coordinates.
+
+    Note:
+        Joints number: K
+
+    Args:
+        joints_3d (np.ndarray([K, 3])): Coordinates of keypoints.
+        rot (float): Rotation angle (degree).
+
+    Returns:
+        joints_3d_rotated
+    """
+    # in-plane rotation
+    # 3D joints are rotated counterclockwise,
+    # so the rot angle is inversed.
+    rot_mat = _construct_rotation_matrix(-rot, 3)
+
+    joints_3d_rotated = np.einsum('ij,kj->ki', rot_mat, joints_3d)
+    joints_3d_rotated = joints_3d_rotated.astype('float32')
+    return joints_3d_rotated
+
+
+def _rotate_smpl_pose(pose, rot):
+    """Rotate SMPL pose parameters. SMPL (https://smpl.is.tue.mpg.de/) is a 3D
+    human model.
+
+    Args:
+        pose (np.ndarray([72])): SMPL pose parameters
+        rot (float): Rotation angle (degree).
+
+    Returns:
+        pose_rotated
+    """
+    pose_rotated = pose.copy()
+    if rot != 0:
+        rot_mat = _construct_rotation_matrix(-rot)
+        orient = pose[:3]
+        # find the rotation of the body in camera frame
+        per_rdg, _ = cv2.Rodrigues(orient)
+        # apply the global rotation to the global orientation
+        res_rot, _ = cv2.Rodrigues(np.dot(rot_mat, per_rdg))
+        pose_rotated[:3] = (res_rot.T)[0]
+
+    return pose_rotated
+
+
+def _flip_joints_3d(joints_3d, joints_3d_visible, flip_pairs):
+    """Flip human joints in 3D space horizontally.
+
+    Note:
+        num_keypoints: K
+
+    Args:
+        joints_3d (np.ndarray([K, 3])): Coordinates of keypoints.
+        joints_3d_visible (np.ndarray([K, 1])): Visibility of keypoints.
+        flip_pairs (list[tuple()]): Pairs of keypoints which are mirrored
+            (for example, left ear -- right ear).
+
+    Returns:
+        joints_3d_flipped, joints_3d_visible_flipped
+    """
+
+    assert len(joints_3d) == len(joints_3d_visible)
+
+    joints_3d_flipped = joints_3d.copy()
+    joints_3d_visible_flipped = joints_3d_visible.copy()
+
+    # Swap left-right parts
+    for left, right in flip_pairs:
+        joints_3d_flipped[left, :] = joints_3d[right, :]
+        joints_3d_flipped[right, :] = joints_3d[left, :]
+
+        joints_3d_visible_flipped[left, :] = joints_3d_visible[right, :]
+        joints_3d_visible_flipped[right, :] = joints_3d_visible[left, :]
+
+    # Flip horizontally
+    joints_3d_flipped[:, 0] = -joints_3d_flipped[:, 0]
+    joints_3d_flipped = joints_3d_flipped * joints_3d_visible_flipped
+
+    return joints_3d_flipped, joints_3d_visible_flipped
+
+
+@PIPELINES.register_module()
+class LoadIUVFromFile:
+    """Loading IUV image from file."""
+
+    def __init__(self, to_float32=False):
+        self.to_float32 = to_float32
+        self.color_type = 'color'
+        # channel relations: iuv->bgr
+        self.channel_order = 'bgr'
+
+    def __call__(self, results):
+        """Loading image from file."""
+        has_iuv = results['has_iuv']
+        use_iuv = results['ann_info']['use_IUV']
+        if has_iuv and use_iuv:
+            iuv_file = results['iuv_file']
+            iuv = mmcv.imread(iuv_file, self.color_type, self.channel_order)
+            if iuv is None:
+                raise ValueError(f'Fail to read {iuv_file}')
+        else:
+            has_iuv = 0
+            iuv = None
+
+        results['has_iuv'] = has_iuv
+        results['iuv'] = iuv
+        return results
+
+
+@PIPELINES.register_module()
+class IUVToTensor:
+    """Transform IUV image to part index mask and uv coordinates image. The 3
+    channels of IUV image means: part index, u coordinates, v coordinates.
+
+    Required key: 'iuv', 'ann_info'.
+    Modifies key: 'part_index', 'uv_coordinates'.
+
+    Args:
+        results (dict): contain all information about training.
+    """
+
+    def __call__(self, results):
+        iuv = results['iuv']
+        if iuv is None:
+            H, W = results['ann_info']['iuv_size']
+            part_index = torch.zeros([1, H, W], dtype=torch.long)
+            uv_coordinates = torch.zeros([2, H, W], dtype=torch.float32)
+        else:
+            part_index = torch.LongTensor(iuv[:, :, 0])[None, :, :]
+            uv_coordinates = torch.FloatTensor(iuv[:, :, 1:]) / 255
+            uv_coordinates = uv_coordinates.permute(2, 0, 1)
+        results['part_index'] = part_index
+        results['uv_coordinates'] = uv_coordinates
+        return results
+
+
+@PIPELINES.register_module()
+class MeshRandomChannelNoise:
+    """Data augmentation with random channel noise.
+
+    Required keys: 'img'
+    Modifies key: 'img'
+
+    Args:
+        noise_factor (float): Multiply each channel with
+         a factor between``[1-scale_factor, 1+scale_factor]``
+    """
+
+    def __init__(self, noise_factor=0.4):
+        self.noise_factor = noise_factor
+
+    def __call__(self, results):
+        """Perform data augmentation with random channel noise."""
+        img = results['img']
+
+        # Each channel is multiplied with a number
+        # in the area [1-self.noise_factor, 1+self.noise_factor]
+        pn = np.random.uniform(1 - self.noise_factor, 1 + self.noise_factor,
+                               (1, 3))
+        img = cv2.multiply(img, pn)
+
+        results['img'] = img
+        return results
+
+
+@PIPELINES.register_module()
+class MeshRandomFlip:
+    """Data augmentation with random image flip.
+
+    Required keys: 'img', 'joints_2d','joints_2d_visible', 'joints_3d',
+    'joints_3d_visible', 'center', 'pose', 'iuv' and 'ann_info'.
+    Modifies key: 'img', 'joints_2d','joints_2d_visible', 'joints_3d',
+    'joints_3d_visible', 'center', 'pose', 'iuv'.
+
+    Args:
+        flip_prob (float): Probability of flip.
+    """
+
+    def __init__(self, flip_prob=0.5):
+        self.flip_prob = flip_prob
+
+    def __call__(self, results):
+        """Perform data augmentation with random image flip."""
+        if np.random.rand() > self.flip_prob:
+            return results
+
+        img = results['img']
+        joints_2d = results['joints_2d']
+        joints_2d_visible = results['joints_2d_visible']
+        joints_3d = results['joints_3d']
+        joints_3d_visible = results['joints_3d_visible']
+        pose = results['pose']
+        center = results['center']
+
+        img = img[:, ::-1, :]
+        pose = _flip_smpl_pose(pose)
+
+        joints_2d, joints_2d_visible = fliplr_joints(
+            joints_2d, joints_2d_visible, img.shape[1],
+            results['ann_info']['flip_pairs'])
+
+        joints_3d, joints_3d_visible = _flip_joints_3d(
+            joints_3d, joints_3d_visible, results['ann_info']['flip_pairs'])
+        center[0] = img.shape[1] - center[0] - 1
+
+        if 'iuv' in results.keys():
+            iuv = results['iuv']
+            if iuv is not None:
+                iuv = _flip_iuv(iuv, results['ann_info']['uv_type'])
+            results['iuv'] = iuv
+
+        results['img'] = img
+        results['joints_2d'] = joints_2d
+        results['joints_2d_visible'] = joints_2d_visible
+        results['joints_3d'] = joints_3d
+        results['joints_3d_visible'] = joints_3d_visible
+        results['pose'] = pose
+        results['center'] = center
+        return results
+
+
+@PIPELINES.register_module()
+class MeshGetRandomScaleRotation:
+    """Data augmentation with random scaling & rotating.
+
+    Required key: 'scale'. Modifies key: 'scale' and 'rotation'.
+
+    Args:
+        rot_factor (int): Rotating to ``[-2*rot_factor, 2*rot_factor]``.
+        scale_factor (float): Scaling to ``[1-scale_factor, 1+scale_factor]``.
+        rot_prob (float): Probability of random rotation.
+    """
+
+    def __init__(self, rot_factor=30, scale_factor=0.25, rot_prob=0.6):
+        self.rot_factor = rot_factor
+        self.scale_factor = scale_factor
+        self.rot_prob = rot_prob
+
+    def __call__(self, results):
+        """Perform data augmentation with random scaling & rotating."""
+        s = results['scale']
+
+        sf = self.scale_factor
+        rf = self.rot_factor
+
+        s_factor = np.clip(np.random.randn() * sf + 1, 1 - sf, 1 + sf)
+        s = s * s_factor
+
+        r_factor = np.clip(np.random.randn() * rf, -rf * 2, rf * 2)
+        r = r_factor if np.random.rand() <= self.rot_prob else 0
+
+        results['scale'] = s
+        results['rotation'] = r
+
+        return results
+
+
+@PIPELINES.register_module()
+class MeshAffine:
+    """Affine transform the image to get input image. Affine transform the 2D
+    keypoints, 3D kepoints and IUV image too.
+
+    Required keys: 'img', 'joints_2d','joints_2d_visible', 'joints_3d',
+    'joints_3d_visible', 'pose', 'iuv', 'ann_info','scale',  'rotation' and
+    'center'. Modifies key: 'img', 'joints_2d','joints_2d_visible',
+    'joints_3d',  'pose', 'iuv'.
+    """
+
+    def __call__(self, results):
+        image_size = results['ann_info']['image_size']
+
+        img = results['img']
+        joints_2d = results['joints_2d']
+        joints_2d_visible = results['joints_2d_visible']
+        joints_3d = results['joints_3d']
+        pose = results['pose']
+
+        c = results['center']
+        s = results['scale']
+        r = results['rotation']
+        trans = get_affine_transform(c, s, r, image_size)
+
+        img = cv2.warpAffine(
+            img,
+            trans, (int(image_size[0]), int(image_size[1])),
+            flags=cv2.INTER_LINEAR)
+
+        for i in range(results['ann_info']['num_joints']):
+            if joints_2d_visible[i, 0] > 0.0:
+                joints_2d[i] = affine_transform(joints_2d[i], trans)
+
+        joints_3d = _rotate_joints_3d(joints_3d, r)
+        pose = _rotate_smpl_pose(pose, r)
+
+        results['img'] = img
+        results['joints_2d'] = joints_2d
+        results['joints_2d_visible'] = joints_2d_visible
+        results['joints_3d'] = joints_3d
+        results['pose'] = pose
+
+        if 'iuv' in results.keys():
+            iuv = results['iuv']
+            if iuv is not None:
+                iuv_size = results['ann_info']['iuv_size']
+                iuv = cv2.warpAffine(
+                    iuv,
+                    trans, (int(iuv_size[0]), int(iuv_size[1])),
+                    flags=cv2.INTER_NEAREST)
+            results['iuv'] = iuv
+
+        return results

third_party/hamer/third-party/ViTPose/mmpose/datasets/pipelines/pose3d_transform.py

@@ -0,0 +1,643 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import copy
+
+import mmcv
+import numpy as np
+import torch
+from mmcv.utils import build_from_cfg
+
+from mmpose.core.camera import CAMERAS
+from mmpose.core.post_processing import fliplr_regression
+from mmpose.datasets.builder import PIPELINES
+
+
+@PIPELINES.register_module()
+class GetRootCenteredPose:
+    """Zero-center the pose around a given root joint. Optionally, the root
+    joint can be removed from the original pose and stored as a separate item.
+
+    Note that the root-centered joints may no longer align with some annotation
+    information (e.g. flip_pairs, num_joints, inference_channel, etc.) due to
+    the removal of the root joint.
+
+    Args:
+        item (str): The name of the pose to apply root-centering.
+        root_index (int): Root joint index in the pose.
+        visible_item (str): The name of the visibility item.
+        remove_root (bool): If true, remove the root joint from the pose
+        root_name (str): Optional. If not none, it will be used as the key to
+            store the root position separated from the original pose.
+
+    Required keys:
+        item
+
+    Modified keys:
+        item, visible_item, root_name
+    """
+
+    def __init__(self,
+                 item,
+                 root_index,
+                 visible_item=None,
+                 remove_root=False,
+                 root_name=None):
+        self.item = item
+        self.root_index = root_index
+        self.remove_root = remove_root
+        self.root_name = root_name
+        self.visible_item = visible_item
+
+    def __call__(self, results):
+        assert self.item in results
+        joints = results[self.item]
+        root_idx = self.root_index
+
+        assert joints.ndim >= 2 and joints.shape[-2] > root_idx,\
+            f'Got invalid joint shape {joints.shape}'
+
+        root = joints[..., root_idx:root_idx + 1, :]
+        joints = joints - root
+
+        results[self.item] = joints
+        if self.root_name is not None:
+            results[self.root_name] = root
+
+        if self.remove_root:
+            results[self.item] = np.delete(
+                results[self.item], root_idx, axis=-2)
+            if self.visible_item is not None:
+                assert self.visible_item in results
+                results[self.visible_item] = np.delete(
+                    results[self.visible_item], root_idx, axis=-2)
+            # Add a flag to avoid latter transforms that rely on the root
+            # joint or the original joint index
+            results[f'{self.item}_root_removed'] = True
+
+            # Save the root index which is necessary to restore the global pose
+            if self.root_name is not None:
+                results[f'{self.root_name}_index'] = self.root_index
+
+        return results
+
+
+@PIPELINES.register_module()
+class NormalizeJointCoordinate:
+    """Normalize the joint coordinate with given mean and std.
+
+    Args:
+        item (str): The name of the pose to normalize.
+        mean (array): Mean values of joint coordinates in shape [K, C].
+        std (array): Std values of joint coordinates in shape [K, C].
+        norm_param_file (str): Optionally load a dict containing `mean` and
+            `std` from a file using `mmcv.load`.
+
+    Required keys:
+        item
+
+    Modified keys:
+        item
+    """
+
+    def __init__(self, item, mean=None, std=None, norm_param_file=None):
+        self.item = item
+        self.norm_param_file = norm_param_file
+        if norm_param_file is not None:
+            norm_param = mmcv.load(norm_param_file)
+            assert 'mean' in norm_param and 'std' in norm_param
+            mean = norm_param['mean']
+            std = norm_param['std']
+        else:
+            assert mean is not None
+            assert std is not None
+
+        self.mean = np.array(mean, dtype=np.float32)
+        self.std = np.array(std, dtype=np.float32)
+
+    def __call__(self, results):
+        assert self.item in results
+        results[self.item] = (results[self.item] - self.mean) / self.std
+        results[f'{self.item}_mean'] = self.mean.copy()
+        results[f'{self.item}_std'] = self.std.copy()
+        return results
+
+
+@PIPELINES.register_module()
+class ImageCoordinateNormalization:
+    """Normalize the 2D joint coordinate with image width and height. Range [0,
+    w] is mapped to [-1, 1], while preserving the aspect ratio.
+
+    Args:
+        item (str|list[str]): The name of the pose to normalize.
+        norm_camera (bool): Whether to normalize camera intrinsics.
+            Default: False.
+        camera_param (dict|None): The camera parameter dict. See the camera
+            class definition for more details. If None is given, the camera
+            parameter will be obtained during processing of each data sample
+            with the key "camera_param".
+
+    Required keys:
+        item
+
+    Modified keys:
+        item (, camera_param)
+    """
+
+    def __init__(self, item, norm_camera=False, camera_param=None):
+        self.item = item
+        if isinstance(self.item, str):
+            self.item = [self.item]
+
+        self.norm_camera = norm_camera
+
+        if camera_param is None:
+            self.static_camera = False
+        else:
+            self.static_camera = True
+            self.camera_param = camera_param
+
+    def __call__(self, results):
+        center = np.array(
+            [0.5 * results['image_width'], 0.5 * results['image_height']],
+            dtype=np.float32)
+        scale = np.array(0.5 * results['image_width'], dtype=np.float32)
+
+        for item in self.item:
+            results[item] = (results[item] - center) / scale
+
+        if self.norm_camera:
+            if self.static_camera:
+                camera_param = copy.deepcopy(self.camera_param)
+            else:
+                assert 'camera_param' in results, \
+                    'Camera parameters are missing.'
+                camera_param = results['camera_param']
+            assert 'f' in camera_param and 'c' in camera_param
+            camera_param['f'] = camera_param['f'] / scale
+            camera_param['c'] = (camera_param['c'] - center[:, None]) / scale
+            if 'camera_param' not in results:
+                results['camera_param'] = dict()
+            results['camera_param'].update(camera_param)
+
+        return results
+
+
+@PIPELINES.register_module()
+class CollectCameraIntrinsics:
+    """Store camera intrinsics in a 1-dim array, including f, c, k, p.
+
+    Args:
+        camera_param (dict|None): The camera parameter dict. See the camera
+            class definition for more details. If None is given, the camera
+            parameter will be obtained during processing of each data sample
+            with the key "camera_param".
+        need_distortion (bool): Whether need distortion parameters k and p.
+            Default: True.
+
+    Required keys:
+        camera_param (if camera parameters are not given in initialization)
+
+    Modified keys:
+        intrinsics
+    """
+
+    def __init__(self, camera_param=None, need_distortion=True):
+        if camera_param is None:
+            self.static_camera = False
+        else:
+            self.static_camera = True
+            self.camera_param = camera_param
+        self.need_distortion = need_distortion
+
+    def __call__(self, results):
+        if self.static_camera:
+            camera_param = copy.deepcopy(self.camera_param)
+        else:
+            assert 'camera_param' in results, 'Camera parameters are missing.'
+            camera_param = results['camera_param']
+        assert 'f' in camera_param and 'c' in camera_param
+        intrinsics = np.concatenate(
+            [camera_param['f'].reshape(2), camera_param['c'].reshape(2)])
+        if self.need_distortion:
+            assert 'k' in camera_param and 'p' in camera_param
+            intrinsics = np.concatenate([
+                intrinsics, camera_param['k'].reshape(3),
+                camera_param['p'].reshape(2)
+            ])
+        results['intrinsics'] = intrinsics
+
+        return results
+
+
+@PIPELINES.register_module()
+class CameraProjection:
+    """Apply camera projection to joint coordinates.
+
+    Args:
+        item (str): The name of the pose to apply camera projection.
+        mode (str): The type of camera projection, supported options are
+
+            - world_to_camera
+            - world_to_pixel
+            - camera_to_world
+            - camera_to_pixel
+        output_name (str|None): The name of the projected pose. If None
+            (default) is given, the projected pose will be stored in place.
+        camera_type (str): The camera class name (should be registered in
+            CAMERA).
+        camera_param (dict|None): The camera parameter dict. See the camera
+            class definition for more details. If None is given, the camera
+            parameter will be obtained during processing of each data sample
+            with the key "camera_param".
+
+    Required keys:
+
+        - item
+        - camera_param (if camera parameters are not given in initialization)
+
+    Modified keys:
+        output_name
+    """
+
+    def __init__(self,
+                 item,
+                 mode,
+                 output_name=None,
+                 camera_type='SimpleCamera',
+                 camera_param=None):
+        self.item = item
+        self.mode = mode
+        self.output_name = output_name
+        self.camera_type = camera_type
+        allowed_mode = {
+            'world_to_camera',
+            'world_to_pixel',
+            'camera_to_world',
+            'camera_to_pixel',
+        }
+        if mode not in allowed_mode:
+            raise ValueError(
+                f'Got invalid mode: {mode}, allowed modes are {allowed_mode}')
+
+        if camera_param is None:
+            self.static_camera = False
+        else:
+            self.static_camera = True
+            self.camera = self._build_camera(camera_param)
+
+    def _build_camera(self, param):
+        cfgs = dict(type=self.camera_type, param=param)
+        return build_from_cfg(cfgs, CAMERAS)
+
+    def __call__(self, results):
+        assert self.item in results
+        joints = results[self.item]
+
+        if self.static_camera:
+            camera = self.camera
+        else:
+            assert 'camera_param' in results, 'Camera parameters are missing.'
+            camera = self._build_camera(results['camera_param'])
+
+        if self.mode == 'world_to_camera':
+            output = camera.world_to_camera(joints)
+        elif self.mode == 'world_to_pixel':
+            output = camera.world_to_pixel(joints)
+        elif self.mode == 'camera_to_world':
+            output = camera.camera_to_world(joints)
+        elif self.mode == 'camera_to_pixel':
+            output = camera.camera_to_pixel(joints)
+        else:
+            raise NotImplementedError
+
+        output_name = self.output_name
+        if output_name is None:
+            output_name = self.item
+
+        results[output_name] = output
+        return results
+
+
+@PIPELINES.register_module()
+class RelativeJointRandomFlip:
+    """Data augmentation with random horizontal joint flip around a root joint.
+
+    Args:
+        item (str|list[str]): The name of the pose to flip.
+        flip_cfg (dict|list[dict]): Configurations of the fliplr_regression
+            function. It should contain the following arguments:
+
+            - ``center_mode``: The mode to set the center location on the \
+                x-axis to flip around.
+            - ``center_x`` or ``center_index``: Set the x-axis location or \
+                the root joint's index to define the flip center.
+
+            Please refer to the docstring of the fliplr_regression function for
+            more details.
+        visible_item (str|list[str]): The name of the visibility item which
+            will be flipped accordingly along with the pose.
+        flip_prob (float): Probability of flip.
+        flip_camera (bool): Whether to flip horizontal distortion coefficients.
+        camera_param (dict|None): The camera parameter dict. See the camera
+            class definition for more details. If None is given, the camera
+            parameter will be obtained during processing of each data sample
+            with the key "camera_param".
+
+    Required keys:
+        item
+
+    Modified keys:
+        item (, camera_param)
+    """
+
+    def __init__(self,
+                 item,
+                 flip_cfg,
+                 visible_item=None,
+                 flip_prob=0.5,
+                 flip_camera=False,
+                 camera_param=None):
+        self.item = item
+        self.flip_cfg = flip_cfg
+        self.vis_item = visible_item
+        self.flip_prob = flip_prob
+        self.flip_camera = flip_camera
+        if camera_param is None:
+            self.static_camera = False
+        else:
+            self.static_camera = True
+            self.camera_param = camera_param
+
+        if isinstance(self.item, str):
+            self.item = [self.item]
+        if isinstance(self.flip_cfg, dict):
+            self.flip_cfg = [self.flip_cfg] * len(self.item)
+        assert len(self.item) == len(self.flip_cfg)
+        if isinstance(self.vis_item, str):
+            self.vis_item = [self.vis_item]
+
+    def __call__(self, results):
+
+        if results.get(f'{self.item}_root_removed', False):
+            raise RuntimeError('The transform RelativeJointRandomFlip should '
+                               f'not be applied to {self.item} whose root '
+                               'joint has been removed and joint indices have '
+                               'been changed')
+
+        if np.random.rand() <= self.flip_prob:
+
+            flip_pairs = results['ann_info']['flip_pairs']
+
+            # flip joint coordinates
+            for i, item in enumerate(self.item):
+                assert item in results
+                joints = results[item]
+
+                joints_flipped = fliplr_regression(joints, flip_pairs,
+                                                   **self.flip_cfg[i])
+
+                results[item] = joints_flipped
+
+            # flip joint visibility
+            for vis_item in self.vis_item:
+                assert vis_item in results
+                visible = results[vis_item]
+                visible_flipped = visible.copy()
+                for left, right in flip_pairs:
+                    visible_flipped[..., left, :] = visible[..., right, :]
+                    visible_flipped[..., right, :] = visible[..., left, :]
+                results[vis_item] = visible_flipped
+
+            # flip horizontal distortion coefficients
+            if self.flip_camera:
+                if self.static_camera:
+                    camera_param = copy.deepcopy(self.camera_param)
+                else:
+                    assert 'camera_param' in results, \
+                        'Camera parameters are missing.'
+                    camera_param = results['camera_param']
+                assert 'c' in camera_param
+                camera_param['c'][0] *= -1
+
+                if 'p' in camera_param:
+                    camera_param['p'][0] *= -1
+
+                if 'camera_param' not in results:
+                    results['camera_param'] = dict()
+                results['camera_param'].update(camera_param)
+
+        return results
+
+
+@PIPELINES.register_module()
+class PoseSequenceToTensor:
+    """Convert pose sequence from numpy array to Tensor.
+
+    The original pose sequence should have a shape of [T,K,C] or [K,C], where
+    T is the sequence length, K and C are keypoint number and dimension. The
+    converted pose sequence will have a shape of [KxC, T].
+
+    Args:
+        item (str): The name of the pose sequence
+
+    Required keys:
+        item
+
+    Modified keys:
+        item
+    """
+
+    def __init__(self, item):
+        self.item = item
+
+    def __call__(self, results):
+        assert self.item in results
+        seq = results[self.item]
+
+        assert isinstance(seq, np.ndarray)
+        assert seq.ndim in {2, 3}
+
+        if seq.ndim == 2:
+            seq = seq[None, ...]
+
+        T = seq.shape[0]
+        seq = seq.transpose(1, 2, 0).reshape(-1, T)
+        results[self.item] = torch.from_numpy(seq)
+
+        return results
+
+
+@PIPELINES.register_module()
+class Generate3DHeatmapTarget:
+    """Generate the target 3d heatmap.
+
+    Required keys: 'joints_3d', 'joints_3d_visible', 'ann_info'.
+    Modified keys: 'target', and 'target_weight'.
+
+    Args:
+        sigma: Sigma of heatmap gaussian.
+        joint_indices (list): Indices of joints used for heatmap generation.
+            If None (default) is given, all joints will be used.
+        max_bound (float): The maximal value of heatmap.
+    """
+
+    def __init__(self, sigma=2, joint_indices=None, max_bound=1.0):
+        self.sigma = sigma
+        self.joint_indices = joint_indices
+        self.max_bound = max_bound
+
+    def __call__(self, results):
+        """Generate the target heatmap."""
+        joints_3d = results['joints_3d']
+        joints_3d_visible = results['joints_3d_visible']
+        cfg = results['ann_info']
+        image_size = cfg['image_size']
+        W, H, D = cfg['heatmap_size']
+        heatmap3d_depth_bound = cfg['heatmap3d_depth_bound']
+        joint_weights = cfg['joint_weights']
+        use_different_joint_weights = cfg['use_different_joint_weights']
+
+        # select the joints used for target generation
+        if self.joint_indices is not None:
+            joints_3d = joints_3d[self.joint_indices, ...]
+            joints_3d_visible = joints_3d_visible[self.joint_indices, ...]
+            joint_weights = joint_weights[self.joint_indices, ...]
+        num_joints = joints_3d.shape[0]
+
+        # get the joint location in heatmap coordinates
+        mu_x = joints_3d[:, 0] * W / image_size[0]
+        mu_y = joints_3d[:, 1] * H / image_size[1]
+        mu_z = (joints_3d[:, 2] / heatmap3d_depth_bound + 0.5) * D
+
+        target = np.zeros([num_joints, D, H, W], dtype=np.float32)
+
+        target_weight = joints_3d_visible[:, 0].astype(np.float32)
+        target_weight = target_weight * (mu_z >= 0) * (mu_z < D)
+        if use_different_joint_weights:
+            target_weight = target_weight * joint_weights
+        target_weight = target_weight[:, None]
+
+        # only compute the voxel value near the joints location
+        tmp_size = 3 * self.sigma
+
+        # get neighboring voxels coordinates
+        x = y = z = np.arange(2 * tmp_size + 1, dtype=np.float32) - tmp_size
+        zz, yy, xx = np.meshgrid(z, y, x)
+        xx = xx[None, ...].astype(np.float32)
+        yy = yy[None, ...].astype(np.float32)
+        zz = zz[None, ...].astype(np.float32)
+        mu_x = mu_x[..., None, None, None]
+        mu_y = mu_y[..., None, None, None]
+        mu_z = mu_z[..., None, None, None]
+        xx, yy, zz = xx + mu_x, yy + mu_y, zz + mu_z
+
+        # round the coordinates
+        xx = xx.round().clip(0, W - 1)
+        yy = yy.round().clip(0, H - 1)
+        zz = zz.round().clip(0, D - 1)
+
+        # compute the target value near joints
+        local_target = \
+            np.exp(-((xx - mu_x)**2 + (yy - mu_y)**2 + (zz - mu_z)**2) /
+                   (2 * self.sigma**2))
+
+        # put the local target value to the full target heatmap
+        local_size = xx.shape[1]
+        idx_joints = np.tile(
+            np.arange(num_joints)[:, None, None, None],
+            [1, local_size, local_size, local_size])
+        idx = np.stack([idx_joints, zz, yy, xx],
+                       axis=-1).astype(int).reshape(-1, 4)
+        target[idx[:, 0], idx[:, 1], idx[:, 2],
+               idx[:, 3]] = local_target.reshape(-1)
+        target = target * self.max_bound
+        results['target'] = target
+        results['target_weight'] = target_weight
+        return results
+
+
+@PIPELINES.register_module()
+class GenerateVoxel3DHeatmapTarget:
+    """Generate the target 3d heatmap.
+
+    Required keys: 'joints_3d', 'joints_3d_visible', 'ann_info_3d'.
+    Modified keys: 'target', and 'target_weight'.
+
+    Args:
+        sigma: Sigma of heatmap gaussian (mm).
+        joint_indices (list): Indices of joints used for heatmap generation.
+            If None (default) is given, all joints will be used.
+    """
+
+    def __init__(self, sigma=200.0, joint_indices=None):
+        self.sigma = sigma  # mm
+        self.joint_indices = joint_indices
+
+    def __call__(self, results):
+        """Generate the target heatmap."""
+        joints_3d = results['joints_3d']
+        joints_3d_visible = results['joints_3d_visible']
+        cfg = results['ann_info']
+
+        num_people = len(joints_3d)
+        num_joints = joints_3d[0].shape[0]
+
+        if self.joint_indices is not None:
+            num_joints = len(self.joint_indices)
+            joint_indices = self.joint_indices
+        else:
+            joint_indices = list(range(num_joints))
+
+        space_size = cfg['space_size']
+        space_center = cfg['space_center']
+        cube_size = cfg['cube_size']
+        grids_x = np.linspace(-space_size[0] / 2, space_size[0] / 2,
+                              cube_size[0]) + space_center[0]
+        grids_y = np.linspace(-space_size[1] / 2, space_size[1] / 2,
+                              cube_size[1]) + space_center[1]
+        grids_z = np.linspace(-space_size[2] / 2, space_size[2] / 2,
+                              cube_size[2]) + space_center[2]
+
+        target = np.zeros(
+            (num_joints, cube_size[0], cube_size[1], cube_size[2]),
+            dtype=np.float32)
+
+        for n in range(num_people):
+            for idx, joint_id in enumerate(joint_indices):
+                mu_x = joints_3d[n][joint_id][0]
+                mu_y = joints_3d[n][joint_id][1]
+                mu_z = joints_3d[n][joint_id][2]
+                vis = joints_3d_visible[n][joint_id][0]
+                if vis < 1:
+                    continue
+                i_x = [
+                    np.searchsorted(grids_x, mu_x - 3 * self.sigma),
+                    np.searchsorted(grids_x, mu_x + 3 * self.sigma, 'right')
+                ]
+                i_y = [
+                    np.searchsorted(grids_y, mu_y - 3 * self.sigma),
+                    np.searchsorted(grids_y, mu_y + 3 * self.sigma, 'right')
+                ]
+                i_z = [
+                    np.searchsorted(grids_z, mu_z - 3 * self.sigma),
+                    np.searchsorted(grids_z, mu_z + 3 * self.sigma, 'right')
+                ]
+                if i_x[0] >= i_x[1] or i_y[0] >= i_y[1] or i_z[0] >= i_z[1]:
+                    continue
+                kernel_xs, kernel_ys, kernel_zs = np.meshgrid(
+                    grids_x[i_x[0]:i_x[1]],
+                    grids_y[i_y[0]:i_y[1]],
+                    grids_z[i_z[0]:i_z[1]],
+                    indexing='ij')
+                g = np.exp(-((kernel_xs - mu_x)**2 + (kernel_ys - mu_y)**2 +
+                             (kernel_zs - mu_z)**2) / (2 * self.sigma**2))
+                target[idx, i_x[0]:i_x[1], i_y[0]:i_y[1], i_z[0]:i_z[1]] \
+                    = np.maximum(target[idx, i_x[0]:i_x[1],
+                                 i_y[0]:i_y[1], i_z[0]:i_z[1]], g)
+
+        target = np.clip(target, 0, 1)
+        if target.shape[0] == 1:
+            target = target[0]
+
+        results['targets_3d'] = target
+
+        return results

third_party/hamer/third-party/ViTPose/mmpose/datasets/pipelines/shared_transform.py

@@ -0,0 +1,527 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import warnings
+from collections.abc import Sequence
+
+import mmcv
+import numpy as np
+from mmcv.parallel import DataContainer as DC
+from mmcv.utils import build_from_cfg
+from numpy import random
+from torchvision.transforms import functional as F
+
+from ..builder import PIPELINES
+
+try:
+    import albumentations
+except ImportError:
+    albumentations = None
+
+
+@PIPELINES.register_module()
+class ToTensor:
+    """Transform image to Tensor.
+
+    Required key: 'img'. Modifies key: 'img'.
+
+    Args:
+        results (dict): contain all information about training.
+    """
+
+    def __call__(self, results):
+        if isinstance(results['img'], (list, tuple)):
+            results['img'] = [F.to_tensor(img) for img in results['img']]
+        else:
+            results['img'] = F.to_tensor(results['img'])
+
+        return results
+
+
+@PIPELINES.register_module()
+class NormalizeTensor:
+    """Normalize the Tensor image (CxHxW), with mean and std.
+
+    Required key: 'img'. Modifies key: 'img'.
+
+    Args:
+        mean (list[float]): Mean values of 3 channels.
+        std (list[float]): Std values of 3 channels.
+    """
+
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, results):
+        if isinstance(results['img'], (list, tuple)):
+            results['img'] = [
+                F.normalize(img, mean=self.mean, std=self.std)
+                for img in results['img']
+            ]
+        else:
+            results['img'] = F.normalize(
+                results['img'], mean=self.mean, std=self.std)
+
+        return results
+
+
+@PIPELINES.register_module()
+class Compose:
+    """Compose a data pipeline with a sequence of transforms.
+
+    Args:
+        transforms (list[dict | callable]): Either config
+          dicts of transforms or transform objects.
+    """
+
+    def __init__(self, transforms):
+        assert isinstance(transforms, Sequence)
+        self.transforms = []
+        for transform in transforms:
+            if isinstance(transform, dict):
+                transform = build_from_cfg(transform, PIPELINES)
+                self.transforms.append(transform)
+            elif callable(transform):
+                self.transforms.append(transform)
+            else:
+                raise TypeError('transform must be callable or a dict, but got'
+                                f' {type(transform)}')
+
+    def __call__(self, data):
+        """Call function to apply transforms sequentially.
+
+        Args:
+            data (dict): A result dict contains the data to transform.
+
+        Returns:
+            dict: Transformed data.
+        """
+        for t in self.transforms:
+            data = t(data)
+            if data is None:
+                return None
+        return data
+
+    def __repr__(self):
+        """Compute the string representation."""
+        format_string = self.__class__.__name__ + '('
+        for t in self.transforms:
+            format_string += f'\n    {t}'
+        format_string += '\n)'
+        return format_string
+
+
+@PIPELINES.register_module()
+class Collect:
+    """Collect data from the loader relevant to the specific task.
+
+    This keeps the items in `keys` as it is, and collect items in `meta_keys`
+    into a meta item called `meta_name`.This is usually the last stage of the
+    data loader pipeline.
+    For example, when keys='imgs', meta_keys=('filename', 'label',
+    'original_shape'), meta_name='img_metas', the results will be a dict with
+    keys 'imgs' and 'img_metas', where 'img_metas' is a DataContainer of
+    another dict with keys 'filename', 'label', 'original_shape'.
+
+    Args:
+        keys (Sequence[str|tuple]): Required keys to be collected. If a tuple
+          (key, key_new) is given as an element, the item retrieved by key will
+          be renamed as key_new in collected data.
+        meta_name (str): The name of the key that contains meta information.
+          This key is always populated. Default: "img_metas".
+        meta_keys (Sequence[str|tuple]): Keys that are collected under
+          meta_name. The contents of the `meta_name` dictionary depends
+          on `meta_keys`.
+    """
+
+    def __init__(self, keys, meta_keys, meta_name='img_metas'):
+        self.keys = keys
+        self.meta_keys = meta_keys
+        self.meta_name = meta_name
+
+    def __call__(self, results):
+        """Performs the Collect formatting.
+
+        Args:
+            results (dict): The resulting dict to be modified and passed
+              to the next transform in pipeline.
+        """
+        if 'ann_info' in results:
+            results.update(results['ann_info'])
+
+        data = {}
+        for key in self.keys:
+            if isinstance(key, tuple):
+                assert len(key) == 2
+                key_src, key_tgt = key[:2]
+            else:
+                key_src = key_tgt = key
+            data[key_tgt] = results[key_src]
+
+        meta = {}
+        if len(self.meta_keys) != 0:
+            for key in self.meta_keys:
+                if isinstance(key, tuple):
+                    assert len(key) == 2
+                    key_src, key_tgt = key[:2]
+                else:
+                    key_src = key_tgt = key
+                meta[key_tgt] = results[key_src]
+        if 'bbox_id' in results:
+            meta['bbox_id'] = results['bbox_id']
+        data[self.meta_name] = DC(meta, cpu_only=True)
+
+        return data
+
+    def __repr__(self):
+        """Compute the string representation."""
+        return (f'{self.__class__.__name__}('
+                f'keys={self.keys}, meta_keys={self.meta_keys})')
+
+
+@PIPELINES.register_module()
+class Albumentation:
+    """Albumentation augmentation (pixel-level transforms only). Adds custom
+    pixel-level transformations from Albumentations library. Please visit
+    `https://albumentations.readthedocs.io` to get more information.
+
+    Note: we only support pixel-level transforms.
+    Please visit `https://github.com/albumentations-team/`
+    `albumentations#pixel-level-transforms`
+    to get more information about pixel-level transforms.
+
+    An example of ``transforms`` is as followed:
+
+    .. code-block:: python
+
+        [
+            dict(
+                type='RandomBrightnessContrast',
+                brightness_limit=[0.1, 0.3],
+                contrast_limit=[0.1, 0.3],
+                p=0.2),
+            dict(type='ChannelShuffle', p=0.1),
+            dict(
+                type='OneOf',
+                transforms=[
+                    dict(type='Blur', blur_limit=3, p=1.0),
+                    dict(type='MedianBlur', blur_limit=3, p=1.0)
+                ],
+                p=0.1),
+        ]
+
+    Args:
+        transforms (list[dict]): A list of Albumentation transformations
+        keymap (dict): Contains {'input key':'albumentation-style key'},
+            e.g., {'img': 'image'}.
+    """
+
+    def __init__(self, transforms, keymap=None):
+        if albumentations is None:
+            raise RuntimeError('albumentations is not installed')
+
+        self.transforms = transforms
+        self.filter_lost_elements = False
+
+        self.aug = albumentations.Compose(
+            [self.albu_builder(t) for t in self.transforms])
+
+        if not keymap:
+            self.keymap_to_albu = {
+                'img': 'image',
+            }
+        else:
+            self.keymap_to_albu = keymap
+        self.keymap_back = {v: k for k, v in self.keymap_to_albu.items()}
+
+    def albu_builder(self, cfg):
+        """Import a module from albumentations.
+
+        It resembles some of :func:`build_from_cfg` logic.
+
+        Args:
+            cfg (dict): Config dict. It should at least contain the key "type".
+
+        Returns:
+            obj: The constructed object.
+        """
+
+        assert isinstance(cfg, dict) and 'type' in cfg
+        args = cfg.copy()
+
+        obj_type = args.pop('type')
+        if mmcv.is_str(obj_type):
+            if albumentations is None:
+                raise RuntimeError('albumentations is not installed')
+            if not hasattr(albumentations.augmentations.transforms, obj_type):
+                warnings.warn('{obj_type} is not pixel-level transformations. '
+                              'Please use with caution.')
+            obj_cls = getattr(albumentations, obj_type)
+        else:
+            raise TypeError(f'type must be a str, but got {type(obj_type)}')
+
+        if 'transforms' in args:
+            args['transforms'] = [
+                self.albu_builder(transform)
+                for transform in args['transforms']
+            ]
+
+        return obj_cls(**args)
+
+    @staticmethod
+    def mapper(d, keymap):
+        """Dictionary mapper.
+
+        Renames keys according to keymap provided.
+
+        Args:
+            d (dict): old dict
+            keymap (dict): {'old_key':'new_key'}
+
+        Returns:
+            dict: new dict.
+        """
+
+        updated_dict = {keymap.get(k, k): v for k, v in d.items()}
+        return updated_dict
+
+    def __call__(self, results):
+        # dict to albumentations format
+        results = self.mapper(results, self.keymap_to_albu)
+
+        results = self.aug(**results)
+        # back to the original format
+        results = self.mapper(results, self.keymap_back)
+
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__ + f'(transforms={self.transforms})'
+        return repr_str
+
+
+@PIPELINES.register_module()
+class PhotometricDistortion:
+    """Apply photometric distortion to image sequentially, every transformation
+    is applied with a probability of 0.5. The position of random contrast is in
+    second or second to last.
+
+    1. random brightness
+    2. random contrast (mode 0)
+    3. convert color from BGR to HSV
+    4. random saturation
+    5. random hue
+    6. convert color from HSV to BGR
+    7. random contrast (mode 1)
+    8. randomly swap channels
+
+    Args:
+        brightness_delta (int): delta of brightness.
+        contrast_range (tuple): range of contrast.
+        saturation_range (tuple): range of saturation.
+        hue_delta (int): delta of hue.
+    """
+
+    def __init__(self,
+                 brightness_delta=32,
+                 contrast_range=(0.5, 1.5),
+                 saturation_range=(0.5, 1.5),
+                 hue_delta=18):
+        self.brightness_delta = brightness_delta
+        self.contrast_lower, self.contrast_upper = contrast_range
+        self.saturation_lower, self.saturation_upper = saturation_range
+        self.hue_delta = hue_delta
+
+    def convert(self, img, alpha=1, beta=0):
+        """Multiple with alpha and add beta with clip."""
+        img = img.astype(np.float32) * alpha + beta
+        img = np.clip(img, 0, 255)
+        return img.astype(np.uint8)
+
+    def brightness(self, img):
+        """Brightness distortion."""
+        if random.randint(2):
+            return self.convert(
+                img,
+                beta=random.uniform(-self.brightness_delta,
+                                    self.brightness_delta))
+        return img
+
+    def contrast(self, img):
+        """Contrast distortion."""
+        if random.randint(2):
+            return self.convert(
+                img,
+                alpha=random.uniform(self.contrast_lower, self.contrast_upper))
+        return img
+
+    def saturation(self, img):
+        # Apply saturation distortion to hsv-formatted img
+        img[:, :, 1] = self.convert(
+            img[:, :, 1],
+            alpha=random.uniform(self.saturation_lower, self.saturation_upper))
+        return img
+
+    def hue(self, img):
+        # Apply hue distortion to hsv-formatted img
+        img[:, :, 0] = (img[:, :, 0].astype(int) +
+                        random.randint(-self.hue_delta, self.hue_delta)) % 180
+        return img
+
+    def swap_channels(self, img):
+        # Apply channel swap
+        if random.randint(2):
+            img = img[..., random.permutation(3)]
+        return img
+
+    def __call__(self, results):
+        """Call function to perform photometric distortion on images.
+
+        Args:
+            results (dict): Result dict from loading pipeline.
+
+        Returns:
+            dict: Result dict with images distorted.
+        """
+
+        img = results['img']
+        # random brightness
+        img = self.brightness(img)
+
+        # mode == 0 --> do random contrast first
+        # mode == 1 --> do random contrast last
+        mode = random.randint(2)
+        if mode == 1:
+            img = self.contrast(img)
+
+        hsv_mode = random.randint(4)
+        if hsv_mode:
+            # random saturation/hue distortion
+            img = mmcv.bgr2hsv(img)
+            if hsv_mode == 1 or hsv_mode == 3:
+                img = self.saturation(img)
+            if hsv_mode == 2 or hsv_mode == 3:
+                img = self.hue(img)
+            img = mmcv.hsv2bgr(img)
+
+        # random contrast
+        if mode == 0:
+            img = self.contrast(img)
+
+        # randomly swap channels
+        self.swap_channels(img)
+
+        results['img'] = img
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        repr_str += (f'(brightness_delta={self.brightness_delta}, '
+                     f'contrast_range=({self.contrast_lower}, '
+                     f'{self.contrast_upper}), '
+                     f'saturation_range=({self.saturation_lower}, '
+                     f'{self.saturation_upper}), '
+                     f'hue_delta={self.hue_delta})')
+        return repr_str
+
+
+@PIPELINES.register_module()
+class MultiItemProcess:
+    """Process each item and merge multi-item results to lists.
+
+    Args:
+        pipeline (dict): Dictionary to construct pipeline for a single item.
+    """
+
+    def __init__(self, pipeline):
+        self.pipeline = Compose(pipeline)
+
+    def __call__(self, results):
+        results_ = {}
+        for idx, result in results.items():
+            single_result = self.pipeline(result)
+            for k, v in single_result.items():
+                if k in results_:
+                    results_[k].append(v)
+                else:
+                    results_[k] = [v]
+
+        return results_
+
+
+@PIPELINES.register_module()
+class DiscardDuplicatedItems:
+
+    def __init__(self, keys_list):
+        """Discard duplicated single-item results.
+
+        Args:
+            keys_list (list): List of keys that need to be deduplicate.
+        """
+        self.keys_list = keys_list
+
+    def __call__(self, results):
+        for k, v in results.items():
+            if k in self.keys_list:
+                assert isinstance(v, Sequence)
+                results[k] = v[0]
+
+        return results
+
+
+@PIPELINES.register_module()
+class MultitaskGatherTarget:
+    """Gather the targets for multitask heads.
+
+    Args:
+        pipeline_list (list[list]): List of pipelines for all heads.
+        pipeline_indices (list[int]): Pipeline index of each head.
+    """
+
+    def __init__(self,
+                 pipeline_list,
+                 pipeline_indices=None,
+                 keys=('target', 'target_weight')):
+        self.keys = keys
+        self.pipelines = []
+        for pipeline in pipeline_list:
+            self.pipelines.append(Compose(pipeline))
+        if pipeline_indices is None:
+            self.pipeline_indices = list(range(len(pipeline_list)))
+        else:
+            self.pipeline_indices = pipeline_indices
+
+    def __call__(self, results):
+        # generate target and target weights using all pipelines
+        pipeline_outputs = []
+        for pipeline in self.pipelines:
+            pipeline_output = pipeline(results)
+            pipeline_outputs.append(pipeline_output.copy())
+
+        for key in self.keys:
+            result_key = []
+            for ind in self.pipeline_indices:
+                result_key.append(pipeline_outputs[ind].get(key, None))
+            results[key] = result_key
+        return results
+
+
+@PIPELINES.register_module()
+class RenameKeys:
+    """Rename the keys.
+
+    Args:
+        key_pairs (Sequence[tuple]): Required keys to be renamed.
+            If a tuple (key_src, key_tgt) is given as an element,
+            the item retrieved by key_src will be renamed as key_tgt.
+    """
+
+    def __init__(self, key_pairs):
+        self.key_pairs = key_pairs
+
+    def __call__(self, results):
+        """Rename keys."""
+        for key_pair in self.key_pairs:
+            assert len(key_pair) == 2
+            key_src, key_tgt = key_pair
+            results[key_tgt] = results.pop(key_src)
+        return results

third_party/hamer/third-party/ViTPose/mmpose/datasets/pipelines/top_down_transform.py

@@ -0,0 +1,736 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import cv2
+import numpy as np
+
+from mmpose.core.post_processing import (affine_transform, fliplr_joints,
+                                         get_affine_transform, get_warp_matrix,
+                                         warp_affine_joints)
+from mmpose.datasets.builder import PIPELINES
+
+
+@PIPELINES.register_module()
+class TopDownRandomFlip:
+    """Data augmentation with random image flip.
+
+    Required keys: 'img', 'joints_3d', 'joints_3d_visible', 'center' and
+    'ann_info'.
+
+    Modifies key: 'img', 'joints_3d', 'joints_3d_visible', 'center' and
+    'flipped'.
+
+    Args:
+        flip (bool): Option to perform random flip.
+        flip_prob (float): Probability of flip.
+    """
+
+    def __init__(self, flip_prob=0.5):
+        self.flip_prob = flip_prob
+
+    def __call__(self, results):
+        """Perform data augmentation with random image flip."""
+        img = results['img']
+        joints_3d = results['joints_3d']
+        joints_3d_visible = results['joints_3d_visible']
+        center = results['center']
+
+        # A flag indicating whether the image is flipped,
+        # which can be used by child class.
+        flipped = False
+        if np.random.rand() <= self.flip_prob:
+            flipped = True
+            if not isinstance(img, list):
+                img = img[:, ::-1, :]
+            else:
+                img = [i[:, ::-1, :] for i in img]
+            if not isinstance(img, list):
+                joints_3d, joints_3d_visible = fliplr_joints(
+                    joints_3d, joints_3d_visible, img.shape[1],
+                    results['ann_info']['flip_pairs'])
+                center[0] = img.shape[1] - center[0] - 1
+            else:
+                joints_3d, joints_3d_visible = fliplr_joints(
+                    joints_3d, joints_3d_visible, img[0].shape[1],
+                    results['ann_info']['flip_pairs'])
+                center[0] = img[0].shape[1] - center[0] - 1
+
+        results['img'] = img
+        results['joints_3d'] = joints_3d
+        results['joints_3d_visible'] = joints_3d_visible
+        results['center'] = center
+        results['flipped'] = flipped
+
+        return results
+
+
+@PIPELINES.register_module()
+class TopDownHalfBodyTransform:
+    """Data augmentation with half-body transform. Keep only the upper body or
+    the lower body at random.
+
+    Required keys: 'joints_3d', 'joints_3d_visible', and 'ann_info'.
+
+    Modifies key: 'scale' and 'center'.
+
+    Args:
+        num_joints_half_body (int): Threshold of performing
+            half-body transform. If the body has fewer number
+            of joints (< num_joints_half_body), ignore this step.
+        prob_half_body (float): Probability of half-body transform.
+    """
+
+    def __init__(self, num_joints_half_body=8, prob_half_body=0.3):
+        self.num_joints_half_body = num_joints_half_body
+        self.prob_half_body = prob_half_body
+
+    @staticmethod
+    def half_body_transform(cfg, joints_3d, joints_3d_visible):
+        """Get center&scale for half-body transform."""
+        upper_joints = []
+        lower_joints = []
+        for joint_id in range(cfg['num_joints']):
+            if joints_3d_visible[joint_id][0] > 0:
+                if joint_id in cfg['upper_body_ids']:
+                    upper_joints.append(joints_3d[joint_id])
+                else:
+                    lower_joints.append(joints_3d[joint_id])
+
+        if np.random.randn() < 0.5 and len(upper_joints) > 2:
+            selected_joints = upper_joints
+        elif len(lower_joints) > 2:
+            selected_joints = lower_joints
+        else:
+            selected_joints = upper_joints
+
+        if len(selected_joints) < 2:
+            return None, None
+
+        selected_joints = np.array(selected_joints, dtype=np.float32)
+        center = selected_joints.mean(axis=0)[:2]
+
+        left_top = np.amin(selected_joints, axis=0)
+
+        right_bottom = np.amax(selected_joints, axis=0)
+
+        w = right_bottom[0] - left_top[0]
+        h = right_bottom[1] - left_top[1]
+
+        aspect_ratio = cfg['image_size'][0] / cfg['image_size'][1]
+
+        if w > aspect_ratio * h:
+            h = w * 1.0 / aspect_ratio
+        elif w < aspect_ratio * h:
+            w = h * aspect_ratio
+
+        scale = np.array([w / 200.0, h / 200.0], dtype=np.float32)
+        scale = scale * 1.5
+        return center, scale
+
+    def __call__(self, results):
+        """Perform data augmentation with half-body transform."""
+        joints_3d = results['joints_3d']
+        joints_3d_visible = results['joints_3d_visible']
+
+        if (np.sum(joints_3d_visible[:, 0]) > self.num_joints_half_body
+                and np.random.rand() < self.prob_half_body):
+
+            c_half_body, s_half_body = self.half_body_transform(
+                results['ann_info'], joints_3d, joints_3d_visible)
+
+            if c_half_body is not None and s_half_body is not None:
+                results['center'] = c_half_body
+                results['scale'] = s_half_body
+
+        return results
+
+
+@PIPELINES.register_module()
+class TopDownGetRandomScaleRotation:
+    """Data augmentation with random scaling & rotating.
+
+    Required key: 'scale'.
+
+    Modifies key: 'scale' and 'rotation'.
+
+    Args:
+        rot_factor (int): Rotating to ``[-2*rot_factor, 2*rot_factor]``.
+        scale_factor (float): Scaling to ``[1-scale_factor, 1+scale_factor]``.
+        rot_prob (float): Probability of random rotation.
+    """
+
+    def __init__(self, rot_factor=40, scale_factor=0.5, rot_prob=0.6):
+        self.rot_factor = rot_factor
+        self.scale_factor = scale_factor
+        self.rot_prob = rot_prob
+
+    def __call__(self, results):
+        """Perform data augmentation with random scaling & rotating."""
+        s = results['scale']
+
+        sf = self.scale_factor
+        rf = self.rot_factor
+
+        s_factor = np.clip(np.random.randn() * sf + 1, 1 - sf, 1 + sf)
+        s = s * s_factor
+
+        r_factor = np.clip(np.random.randn() * rf, -rf * 2, rf * 2)
+        r = r_factor if np.random.rand() <= self.rot_prob else 0
+
+        results['scale'] = s
+        results['rotation'] = r
+
+        return results
+
+
+@PIPELINES.register_module()
+class TopDownAffine:
+    """Affine transform the image to make input.
+
+    Required keys:'img', 'joints_3d', 'joints_3d_visible', 'ann_info','scale',
+    'rotation' and 'center'.
+
+    Modified keys:'img', 'joints_3d', and 'joints_3d_visible'.
+
+    Args:
+        use_udp (bool): To use unbiased data processing.
+            Paper ref: Huang et al. The Devil is in the Details: Delving into
+            Unbiased Data Processing for Human Pose Estimation (CVPR 2020).
+    """
+
+    def __init__(self, use_udp=False):
+        self.use_udp = use_udp
+
+    def __call__(self, results):
+        image_size = results['ann_info']['image_size']
+
+        img = results['img']
+        joints_3d = results['joints_3d']
+        joints_3d_visible = results['joints_3d_visible']
+        c = results['center']
+        s = results['scale']
+        r = results['rotation']
+
+        if self.use_udp:
+            trans = get_warp_matrix(r, c * 2.0, image_size - 1.0, s * 200.0)
+            if not isinstance(img, list):
+                img = cv2.warpAffine(
+                    img,
+                    trans, (int(image_size[0]), int(image_size[1])),
+                    flags=cv2.INTER_LINEAR)
+            else:
+                img = [
+                    cv2.warpAffine(
+                        i,
+                        trans, (int(image_size[0]), int(image_size[1])),
+                        flags=cv2.INTER_LINEAR) for i in img
+                ]
+
+            joints_3d[:, 0:2] = \
+                warp_affine_joints(joints_3d[:, 0:2].copy(), trans)
+
+        else:
+            trans = get_affine_transform(c, s, r, image_size)
+            if not isinstance(img, list):
+                img = cv2.warpAffine(
+                    img,
+                    trans, (int(image_size[0]), int(image_size[1])),
+                    flags=cv2.INTER_LINEAR)
+            else:
+                img = [
+                    cv2.warpAffine(
+                        i,
+                        trans, (int(image_size[0]), int(image_size[1])),
+                        flags=cv2.INTER_LINEAR) for i in img
+                ]
+            for i in range(results['ann_info']['num_joints']):
+                if joints_3d_visible[i, 0] > 0.0:
+                    joints_3d[i,
+                              0:2] = affine_transform(joints_3d[i, 0:2], trans)
+
+        results['img'] = img
+        results['joints_3d'] = joints_3d
+        results['joints_3d_visible'] = joints_3d_visible
+
+        return results
+
+
+@PIPELINES.register_module()
+class TopDownGenerateTarget:
+    """Generate the target heatmap.
+
+    Required keys: 'joints_3d', 'joints_3d_visible', 'ann_info'.
+
+    Modified keys: 'target', and 'target_weight'.
+
+    Args:
+        sigma: Sigma of heatmap gaussian for 'MSRA' approach.
+        kernel: Kernel of heatmap gaussian for 'Megvii' approach.
+        encoding (str): Approach to generate target heatmaps.
+            Currently supported approaches: 'MSRA', 'Megvii', 'UDP'.
+            Default:'MSRA'
+        unbiased_encoding (bool): Option to use unbiased
+            encoding methods.
+            Paper ref: Zhang et al. Distribution-Aware Coordinate
+            Representation for Human Pose Estimation (CVPR 2020).
+        keypoint_pose_distance: Keypoint pose distance for UDP.
+            Paper ref: Huang et al. The Devil is in the Details: Delving into
+            Unbiased Data Processing for Human Pose Estimation (CVPR 2020).
+        target_type (str): supported targets: 'GaussianHeatmap',
+            'CombinedTarget'. Default:'GaussianHeatmap'
+            CombinedTarget: The combination of classification target
+            (response map) and regression target (offset map).
+            Paper ref: Huang et al. The Devil is in the Details: Delving into
+            Unbiased Data Processing for Human Pose Estimation (CVPR 2020).
+    """
+
+    def __init__(self,
+                 sigma=2,
+                 kernel=(11, 11),
+                 valid_radius_factor=0.0546875,
+                 target_type='GaussianHeatmap',
+                 encoding='MSRA',
+                 unbiased_encoding=False):
+        self.sigma = sigma
+        self.unbiased_encoding = unbiased_encoding
+        self.kernel = kernel
+        self.valid_radius_factor = valid_radius_factor
+        self.target_type = target_type
+        self.encoding = encoding
+
+    def _msra_generate_target(self, cfg, joints_3d, joints_3d_visible, sigma):
+        """Generate the target heatmap via "MSRA" approach.
+
+        Args:
+            cfg (dict): data config
+            joints_3d: np.ndarray ([num_joints, 3])
+            joints_3d_visible: np.ndarray ([num_joints, 3])
+            sigma: Sigma of heatmap gaussian
+        Returns:
+            tuple: A tuple containing targets.
+
+            - target: Target heatmaps.
+            - target_weight: (1: visible, 0: invisible)
+        """
+        num_joints = cfg['num_joints']
+        image_size = cfg['image_size']
+        W, H = cfg['heatmap_size']
+        joint_weights = cfg['joint_weights']
+        use_different_joint_weights = cfg['use_different_joint_weights']
+
+        target_weight = np.zeros((num_joints, 1), dtype=np.float32)
+        target = np.zeros((num_joints, H, W), dtype=np.float32)
+
+        # 3-sigma rule
+        tmp_size = sigma * 3
+
+        if self.unbiased_encoding:
+            for joint_id in range(num_joints):
+                target_weight[joint_id] = joints_3d_visible[joint_id, 0]
+
+                feat_stride = image_size / [W, H]
+                mu_x = joints_3d[joint_id][0] / feat_stride[0]
+                mu_y = joints_3d[joint_id][1] / feat_stride[1]
+                # Check that any part of the gaussian is in-bounds
+                ul = [mu_x - tmp_size, mu_y - tmp_size]
+                br = [mu_x + tmp_size + 1, mu_y + tmp_size + 1]
+                if ul[0] >= W or ul[1] >= H or br[0] < 0 or br[1] < 0:
+                    target_weight[joint_id] = 0
+
+                if target_weight[joint_id] == 0:
+                    continue
+
+                x = np.arange(0, W, 1, np.float32)
+                y = np.arange(0, H, 1, np.float32)
+                y = y[:, None]
+
+                if target_weight[joint_id] > 0.5:
+                    target[joint_id] = np.exp(-((x - mu_x)**2 +
+                                                (y - mu_y)**2) /
+                                              (2 * sigma**2))
+        else:
+            for joint_id in range(num_joints):
+                target_weight[joint_id] = joints_3d_visible[joint_id, 0]
+
+                feat_stride = image_size / [W, H]
+                mu_x = int(joints_3d[joint_id][0] / feat_stride[0] + 0.5)
+                mu_y = int(joints_3d[joint_id][1] / feat_stride[1] + 0.5)
+                # Check that any part of the gaussian is in-bounds
+                ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]
+                br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]
+                if ul[0] >= W or ul[1] >= H or br[0] < 0 or br[1] < 0:
+                    target_weight[joint_id] = 0
+
+                if target_weight[joint_id] > 0.5:
+                    size = 2 * tmp_size + 1
+                    x = np.arange(0, size, 1, np.float32)
+                    y = x[:, None]
+                    x0 = y0 = size // 2
+                    # The gaussian is not normalized,
+                    # we want the center value to equal 1
+                    g = np.exp(-((x - x0)**2 + (y - y0)**2) / (2 * sigma**2))
+
+                    # Usable gaussian range
+                    g_x = max(0, -ul[0]), min(br[0], W) - ul[0]
+                    g_y = max(0, -ul[1]), min(br[1], H) - ul[1]
+                    # Image range
+                    img_x = max(0, ul[0]), min(br[0], W)
+                    img_y = max(0, ul[1]), min(br[1], H)
+
+                    target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] = \
+                        g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
+
+        if use_different_joint_weights:
+            target_weight = np.multiply(target_weight, joint_weights)
+
+        return target, target_weight
+
+    def _megvii_generate_target(self, cfg, joints_3d, joints_3d_visible,
+                                kernel):
+        """Generate the target heatmap via "Megvii" approach.
+
+        Args:
+            cfg (dict): data config
+            joints_3d: np.ndarray ([num_joints, 3])
+            joints_3d_visible: np.ndarray ([num_joints, 3])
+            kernel: Kernel of heatmap gaussian
+
+        Returns:
+            tuple: A tuple containing targets.
+
+            - target: Target heatmaps.
+            - target_weight: (1: visible, 0: invisible)
+        """
+
+        num_joints = cfg['num_joints']
+        image_size = cfg['image_size']
+        W, H = cfg['heatmap_size']
+        heatmaps = np.zeros((num_joints, H, W), dtype='float32')
+        target_weight = np.zeros((num_joints, 1), dtype=np.float32)
+
+        for i in range(num_joints):
+            target_weight[i] = joints_3d_visible[i, 0]
+
+            if target_weight[i] < 1:
+                continue
+
+            target_y = int(joints_3d[i, 1] * H / image_size[1])
+            target_x = int(joints_3d[i, 0] * W / image_size[0])
+
+            if (target_x >= W or target_x < 0) \
+                    or (target_y >= H or target_y < 0):
+                target_weight[i] = 0
+                continue
+
+            heatmaps[i, target_y, target_x] = 1
+            heatmaps[i] = cv2.GaussianBlur(heatmaps[i], kernel, 0)
+            maxi = heatmaps[i, target_y, target_x]
+
+            heatmaps[i] /= maxi / 255
+
+        return heatmaps, target_weight
+
+    def _udp_generate_target(self, cfg, joints_3d, joints_3d_visible, factor,
+                             target_type):
+        """Generate the target heatmap via 'UDP' approach. Paper ref: Huang et
+        al. The Devil is in the Details: Delving into Unbiased Data Processing
+        for Human Pose Estimation (CVPR 2020).
+
+        Note:
+            - num keypoints: K
+            - heatmap height: H
+            - heatmap width: W
+            - num target channels: C
+            - C = K if target_type=='GaussianHeatmap'
+            - C = 3*K if target_type=='CombinedTarget'
+
+        Args:
+            cfg (dict): data config
+            joints_3d (np.ndarray[K, 3]): Annotated keypoints.
+            joints_3d_visible (np.ndarray[K, 3]): Visibility of keypoints.
+            factor (float): kernel factor for GaussianHeatmap target or
+                valid radius factor for CombinedTarget.
+            target_type (str): 'GaussianHeatmap' or 'CombinedTarget'.
+                GaussianHeatmap: Heatmap target with gaussian distribution.
+                CombinedTarget: The combination of classification target
+                (response map) and regression target (offset map).
+
+        Returns:
+            tuple: A tuple containing targets.
+
+            - target (np.ndarray[C, H, W]): Target heatmaps.
+            - target_weight (np.ndarray[K, 1]): (1: visible, 0: invisible)
+        """
+        num_joints = cfg['num_joints']
+        image_size = cfg['image_size']
+        heatmap_size = cfg['heatmap_size']
+        joint_weights = cfg['joint_weights']
+        use_different_joint_weights = cfg['use_different_joint_weights']
+
+        target_weight = np.ones((num_joints, 1), dtype=np.float32)
+        target_weight[:, 0] = joints_3d_visible[:, 0]
+
+        if target_type.lower() == 'GaussianHeatmap'.lower():
+            target = np.zeros((num_joints, heatmap_size[1], heatmap_size[0]),
+                              dtype=np.float32)
+
+            tmp_size = factor * 3
+
+            # prepare for gaussian
+            size = 2 * tmp_size + 1
+            x = np.arange(0, size, 1, np.float32)
+            y = x[:, None]
+
+            for joint_id in range(num_joints):
+                feat_stride = (image_size - 1.0) / (heatmap_size - 1.0)
+                mu_x = int(joints_3d[joint_id][0] / feat_stride[0] + 0.5)
+                mu_y = int(joints_3d[joint_id][1] / feat_stride[1] + 0.5)
+                # Check that any part of the gaussian is in-bounds
+                ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]
+                br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]
+                if ul[0] >= heatmap_size[0] or ul[1] >= heatmap_size[1] \
+                        or br[0] < 0 or br[1] < 0:
+                    # If not, just return the image as is
+                    target_weight[joint_id] = 0
+                    continue
+
+                # # Generate gaussian
+                mu_x_ac = joints_3d[joint_id][0] / feat_stride[0]
+                mu_y_ac = joints_3d[joint_id][1] / feat_stride[1]
+                x0 = y0 = size // 2
+                x0 += mu_x_ac - mu_x
+                y0 += mu_y_ac - mu_y
+                g = np.exp(-((x - x0)**2 + (y - y0)**2) / (2 * factor**2))
+
+                # Usable gaussian range
+                g_x = max(0, -ul[0]), min(br[0], heatmap_size[0]) - ul[0]
+                g_y = max(0, -ul[1]), min(br[1], heatmap_size[1]) - ul[1]
+                # Image range
+                img_x = max(0, ul[0]), min(br[0], heatmap_size[0])
+                img_y = max(0, ul[1]), min(br[1], heatmap_size[1])
+
+                v = target_weight[joint_id]
+                if v > 0.5:
+                    target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] = \
+                        g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
+
+        elif target_type.lower() == 'CombinedTarget'.lower():
+            target = np.zeros(
+                (num_joints, 3, heatmap_size[1] * heatmap_size[0]),
+                dtype=np.float32)
+            feat_width = heatmap_size[0]
+            feat_height = heatmap_size[1]
+            feat_x_int = np.arange(0, feat_width)
+            feat_y_int = np.arange(0, feat_height)
+            feat_x_int, feat_y_int = np.meshgrid(feat_x_int, feat_y_int)
+            feat_x_int = feat_x_int.flatten()
+            feat_y_int = feat_y_int.flatten()
+            # Calculate the radius of the positive area in classification
+            #   heatmap.
+            valid_radius = factor * heatmap_size[1]
+            feat_stride = (image_size - 1.0) / (heatmap_size - 1.0)
+            for joint_id in range(num_joints):
+                mu_x = joints_3d[joint_id][0] / feat_stride[0]
+                mu_y = joints_3d[joint_id][1] / feat_stride[1]
+                x_offset = (mu_x - feat_x_int) / valid_radius
+                y_offset = (mu_y - feat_y_int) / valid_radius
+                dis = x_offset**2 + y_offset**2
+                keep_pos = np.where(dis <= 1)[0]
+                v = target_weight[joint_id]
+                if v > 0.5:
+                    target[joint_id, 0, keep_pos] = 1
+                    target[joint_id, 1, keep_pos] = x_offset[keep_pos]
+                    target[joint_id, 2, keep_pos] = y_offset[keep_pos]
+            target = target.reshape(num_joints * 3, heatmap_size[1],
+                                    heatmap_size[0])
+        else:
+            raise ValueError('target_type should be either '
+                             "'GaussianHeatmap' or 'CombinedTarget'")
+
+        if use_different_joint_weights:
+            target_weight = np.multiply(target_weight, joint_weights)
+
+        return target, target_weight
+
+    def __call__(self, results):
+        """Generate the target heatmap."""
+        joints_3d = results['joints_3d']
+        joints_3d_visible = results['joints_3d_visible']
+
+        assert self.encoding in ['MSRA', 'Megvii', 'UDP']
+
+        if self.encoding == 'MSRA':
+            if isinstance(self.sigma, list):
+                num_sigmas = len(self.sigma)
+                cfg = results['ann_info']
+                num_joints = cfg['num_joints']
+                heatmap_size = cfg['heatmap_size']
+
+                target = np.empty(
+                    (0, num_joints, heatmap_size[1], heatmap_size[0]),
+                    dtype=np.float32)
+                target_weight = np.empty((0, num_joints, 1), dtype=np.float32)
+                for i in range(num_sigmas):
+                    target_i, target_weight_i = self._msra_generate_target(
+                        cfg, joints_3d, joints_3d_visible, self.sigma[i])
+                    target = np.concatenate([target, target_i[None]], axis=0)
+                    target_weight = np.concatenate(
+                        [target_weight, target_weight_i[None]], axis=0)
+            else:
+                target, target_weight = self._msra_generate_target(
+                    results['ann_info'], joints_3d, joints_3d_visible,
+                    self.sigma)
+
+        elif self.encoding == 'Megvii':
+            if isinstance(self.kernel, list):
+                num_kernels = len(self.kernel)
+                cfg = results['ann_info']
+                num_joints = cfg['num_joints']
+                W, H = cfg['heatmap_size']
+
+                target = np.empty((0, num_joints, H, W), dtype=np.float32)
+                target_weight = np.empty((0, num_joints, 1), dtype=np.float32)
+                for i in range(num_kernels):
+                    target_i, target_weight_i = self._megvii_generate_target(
+                        cfg, joints_3d, joints_3d_visible, self.kernel[i])
+                    target = np.concatenate([target, target_i[None]], axis=0)
+                    target_weight = np.concatenate(
+                        [target_weight, target_weight_i[None]], axis=0)
+            else:
+                target, target_weight = self._megvii_generate_target(
+                    results['ann_info'], joints_3d, joints_3d_visible,
+                    self.kernel)
+
+        elif self.encoding == 'UDP':
+            if self.target_type.lower() == 'CombinedTarget'.lower():
+                factors = self.valid_radius_factor
+                channel_factor = 3
+            elif self.target_type.lower() == 'GaussianHeatmap'.lower():
+                factors = self.sigma
+                channel_factor = 1
+            else:
+                raise ValueError('target_type should be either '
+                                 "'GaussianHeatmap' or 'CombinedTarget'")
+            if isinstance(factors, list):
+                num_factors = len(factors)
+                cfg = results['ann_info']
+                num_joints = cfg['num_joints']
+                W, H = cfg['heatmap_size']
+
+                target = np.empty((0, channel_factor * num_joints, H, W),
+                                  dtype=np.float32)
+                target_weight = np.empty((0, num_joints, 1), dtype=np.float32)
+                for i in range(num_factors):
+                    target_i, target_weight_i = self._udp_generate_target(
+                        cfg, joints_3d, joints_3d_visible, factors[i],
+                        self.target_type)
+                    target = np.concatenate([target, target_i[None]], axis=0)
+                    target_weight = np.concatenate(
+                        [target_weight, target_weight_i[None]], axis=0)
+            else:
+                target, target_weight = self._udp_generate_target(
+                    results['ann_info'], joints_3d, joints_3d_visible, factors,
+                    self.target_type)
+        else:
+            raise ValueError(
+                f'Encoding approach {self.encoding} is not supported!')
+
+        if results['ann_info'].get('max_num_joints', None) is not None:
+            W, H = results['ann_info']['heatmap_size']
+            padded_length = int(results['ann_info'].get('max_num_joints') - results['ann_info'].get('num_joints'))
+            target_weight = np.concatenate([target_weight, np.zeros((padded_length, 1), dtype=np.float32)], 0)
+            target = np.concatenate([target, np.zeros((padded_length, H, W), dtype=np.float32)], 0)
+
+        results['target'] = target
+        results['target_weight'] = target_weight
+
+        results['dataset_idx'] = results['ann_info'].get('dataset_idx', 0)
+
+        return results
+
+
+@PIPELINES.register_module()
+class TopDownGenerateTargetRegression:
+    """Generate the target regression vector (coordinates).
+
+    Required keys: 'joints_3d', 'joints_3d_visible', 'ann_info'. Modified keys:
+    'target', and 'target_weight'.
+    """
+
+    def __init__(self):
+        pass
+
+    def _generate_target(self, cfg, joints_3d, joints_3d_visible):
+        """Generate the target regression vector.
+
+        Args:
+            cfg (dict): data config
+            joints_3d: np.ndarray([num_joints, 3])
+            joints_3d_visible: np.ndarray([num_joints, 3])
+
+        Returns:
+             target, target_weight(1: visible, 0: invisible)
+        """
+        image_size = cfg['image_size']
+        joint_weights = cfg['joint_weights']
+        use_different_joint_weights = cfg['use_different_joint_weights']
+
+        mask = (joints_3d[:, 0] >= 0) * (
+            joints_3d[:, 0] <= image_size[0] - 1) * (joints_3d[:, 1] >= 0) * (
+                joints_3d[:, 1] <= image_size[1] - 1)
+
+        target = joints_3d[:, :2] / image_size
+
+        target = target.astype(np.float32)
+        target_weight = joints_3d_visible[:, :2] * mask[:, None]
+
+        if use_different_joint_weights:
+            target_weight = np.multiply(target_weight, joint_weights)
+
+        return target, target_weight
+
+    def __call__(self, results):
+        """Generate the target heatmap."""
+        joints_3d = results['joints_3d']
+        joints_3d_visible = results['joints_3d_visible']
+
+        target, target_weight = self._generate_target(results['ann_info'],
+                                                      joints_3d,
+                                                      joints_3d_visible)
+
+        results['target'] = target
+        results['target_weight'] = target_weight
+
+        return results
+
+
+@PIPELINES.register_module()
+class TopDownRandomTranslation:
+    """Data augmentation with random translation.
+
+    Required key: 'scale' and 'center'.
+
+    Modifies key: 'center'.
+
+    Note:
+        - bbox height: H
+        - bbox width: W
+
+    Args:
+        trans_factor (float): Translating center to
+            ``[-trans_factor, trans_factor] * [W, H] + center``.
+        trans_prob (float): Probability of random translation.
+    """
+
+    def __init__(self, trans_factor=0.15, trans_prob=1.0):
+        self.trans_factor = trans_factor
+        self.trans_prob = trans_prob
+
+    def __call__(self, results):
+        """Perform data augmentation with random translation."""
+        center = results['center']
+        scale = results['scale']
+        if np.random.rand() <= self.trans_prob:
+            # reference bbox size is [200, 200] pixels
+            center += self.trans_factor * np.random.uniform(
+                -1, 1, size=2) * scale * 200
+        results['center'] = center
+        return results