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-
| 254 | |
| 255 | |
| 256 | class TopDownHalfBodyTransform: |
| 257 | """Data augmentation with half-body transform. Keep only the upper body or |
| 258 | the lower body at random. |
| 259 | |
| 260 | Required keys: 'joints_3d', 'joints_3d_visible', and 'ann_info'. |
| 261 | Modifies key: 'scale' and 'center'. |
| 262 | |
| 263 | Args: |
| 264 | num_joints_half_body (int): Threshold of performing |
| 265 | half-body transform. If the body has fewer number |
| 266 | of joints (< num_joints_half_body), ignore this step. |
| 267 | prob_half_body (float): Probability of half-body transform. |
| 268 | """ |
| 269 | |
| 270 | def __init__(self, num_joints_half_body=8, prob_half_body=0.3): |
| 271 | self.num_joints_half_body = num_joints_half_body |
| 272 | self.prob_half_body = prob_half_body |
| 273 | |
| 274 | @staticmethod |
| 275 | def half_body_transform(cfg, joints_3d, joints_3d_visible): |
| 276 | """Get center&scale for half-body transform.""" |
| 277 | upper_joints = [] |
| 278 | lower_joints = [] |
| 279 | for joint_id in range(cfg['num_joints']): |
| 280 | if joints_3d_visible[joint_id][0] > 0: |
| 281 | if joint_id in cfg['upper_body_ids']: |
| 282 | upper_joints.append(joints_3d[joint_id]) |
| 283 | else: |
| 284 | lower_joints.append(joints_3d[joint_id]) |
| 285 | |
| 286 | if np.random.randn() < 0.5 and len(upper_joints) > 2: |
| 287 | selected_joints = upper_joints |
| 288 | elif len(lower_joints) > 2: |
| 289 | selected_joints = lower_joints |
| 290 | else: |
| 291 | selected_joints = upper_joints |
| 292 | |
| 293 | if len(selected_joints) < 2: |
| 294 | return None, None |
| 295 | |
| 296 | selected_joints = np.array(selected_joints, dtype=np.float32) |
| 297 | center = selected_joints.mean(axis=0)[:2] |
| 298 | |
| 299 | left_top = np.amin(selected_joints, axis=0) |
| 300 | |
| 301 | right_bottom = np.amax(selected_joints, axis=0) |
| 302 | |
| 303 | w = right_bottom[0] - left_top[0] |
| 304 | h = right_bottom[1] - left_top[1] |
| 305 | |
| 306 | aspect_ratio = cfg['image_size'][0] / cfg['image_size'][1] |
| 307 | |
| 308 | if w > aspect_ratio * h: |
| 309 | h = w * 1.0 / aspect_ratio |
| 310 | elif w < aspect_ratio * h: |
| 311 | w = h * aspect_ratio |
| 312 | |
| 313 | scale = np.array([w / 200.0, h / 200.0], dtype=np.float32) |