Spaces:
Running
on
Zero
Running
on
Zero
File size: 6,063 Bytes
a249588 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Dict, Optional, Tuple
import cv2
import numpy as np
from mmcv.transforms import BaseTransform
from mmengine import is_seq_of
from mmpose.registry import TRANSFORMS
from mmpose.structures.bbox import get_udp_warp_matrix, get_warp_matrix, bbox_cs2xyxy, bbox_xyxy2cs
@TRANSFORMS.register_module()
class TopdownAffine(BaseTransform):
"""Get the bbox image as the model input by affine transform.
Required Keys:
- img
- bbox_center
- bbox_scale
- bbox_rotation (optional)
- keypoints (optional)
Modified Keys:
- img
- bbox_scale
Added Keys:
- input_size
- transformed_keypoints
- bbox_mask
Args:
input_size (Tuple[int, int]): The input image size of the model in
[w, h]. The bbox region will be cropped and resize to `input_size`
use_udp (bool): Whether use unbiased data processing. See
`UDP (CVPR 2020)`_ for details. Defaults to ``False``
.. _`UDP (CVPR 2020)`: https://arxiv.org/abs/1911.07524
"""
def __init__(self,
input_size: Tuple[int, int],
input_padding: float = 1.25,
use_udp: bool = False) -> None:
super().__init__()
assert is_seq_of(input_size, int) and len(input_size) == 2, (
f'Invalid input_size {input_size}')
self.input_size = input_size
self.use_udp = use_udp
self.input_padding = input_padding
@staticmethod
def _fix_aspect_ratio(bbox_scale: np.ndarray, aspect_ratio: float):
"""Reshape the bbox to a fixed aspect ratio.
Args:
bbox_scale (np.ndarray): The bbox scales (w, h) in shape (n, 2)
aspect_ratio (float): The ratio of ``w/h``
Returns:
np.darray: The reshaped bbox scales in (n, 2)
"""
w, h = np.hsplit(bbox_scale, [1])
bbox_scale = np.where(w > h * aspect_ratio,
np.hstack([w, w / aspect_ratio]),
np.hstack([h * aspect_ratio, h]))
return bbox_scale
def transform(self, results: Dict) -> Optional[dict]:
"""The transform function of :class:`TopdownAffine`.
See ``transform()`` method of :class:`BaseTransform` for details.
Args:
results (dict): The result dict
Returns:
dict: The result dict.
"""
w, h = self.input_size
warp_size = (int(w), int(h))
img_h, img_w = results['img'].shape[:2]
bbox_xyxy = results['bbox_xyxy_wrt_input'].flatten()
bbox_xyxy[:2] = np.maximum(bbox_xyxy[:2], 0)
bbox_xyxy[2:4] = np.minimum(bbox_xyxy[2:4], [img_w, img_h])
x0, y0, x1, y1 = bbox_xyxy[:4].astype(int)
bbox_mask = np.zeros((img_h, img_w), dtype=np.uint8)
bbox_mask[y0:y1, x0:x1] = 1
# Take the bbox wrt the input
bbox_xyxy_wrt_input = results.get('bbox_xyxy_wrt_input', None)
if bbox_xyxy_wrt_input is not None:
_c, _s = bbox_xyxy2cs(bbox_xyxy_wrt_input, padding=self.input_padding)
results['bbox_center'] = _c.reshape(1, 2)
results['bbox_scale'] = _s.reshape(1, 2)
# reshape bbox to fixed aspect ratio
results['bbox_scale'] = self._fix_aspect_ratio(
results['bbox_scale'], aspect_ratio=w / h)
# TODO: support multi-instance
assert results['bbox_center'].shape[0] == 1, (
'Top-down heatmap only supports single instance. Got invalid '
f'shape of bbox_center {results["bbox_center"].shape}.')
center = results['bbox_center'][0]
scale = results['bbox_scale'][0]
if 'bbox_rotation' in results:
rot = results['bbox_rotation'][0]
else:
rot = 0.
if self.use_udp:
warp_mat = get_udp_warp_matrix(
center, scale, rot, output_size=(w, h))
else:
warp_mat = get_warp_matrix(center, scale, rot, output_size=(w, h))
if isinstance(results['img'], list):
results['img'] = [
cv2.warpAffine(
img, warp_mat, warp_size, flags=cv2.INTER_LINEAR)
for img in results['img']
]
else:
results['img'] = cv2.warpAffine(
results['img'], warp_mat, warp_size, flags=cv2.INTER_LINEAR)
bbox_mask = cv2.warpAffine(
bbox_mask, warp_mat, warp_size, flags=cv2.INTER_LINEAR)
bbox_mask = bbox_mask.reshape(1, h, w)
results['bbox_mask'] = bbox_mask
if results.get('keypoints', None) is not None:
if results.get('transformed_keypoints', None) is not None:
transformed_keypoints = results['transformed_keypoints'].copy()
else:
transformed_keypoints = results['keypoints'].copy()
# Only transform (x, y) coordinates
transformed_keypoints[..., :2] = cv2.transform(
transformed_keypoints[..., :2], warp_mat)
results['transformed_keypoints'] = transformed_keypoints
if results.get('bbox_xyxy_wrt_input', None) is not None:
bbox_xyxy_wrt_input = results['bbox_xyxy_wrt_input'].copy()
bbox_xyxy_wrt_input = bbox_xyxy_wrt_input.reshape(1, 2, 2)
bbox_xyxy_wrt_input = cv2.transform(
bbox_xyxy_wrt_input, warp_mat)
results['bbox_xyxy_wrt_input'] = bbox_xyxy_wrt_input.reshape(1, 4)
results['input_size'] = (w, h)
results['input_center'] = center
results['input_scale'] = scale
return results
def __repr__(self) -> str:
"""print the basic information of the transform.
Returns:
str: Formatted string.
"""
repr_str = self.__class__.__name__
repr_str += f'(input_size={self.input_size}, '
repr_str += f'use_udp={self.use_udp})'
return repr_str
|