You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
113 lines
4.1 KiB
113 lines
4.1 KiB
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 TopDownAffineSam:
|
|
"""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']
|
|
# 修改
|
|
sam_image_size = results['ann_info']['sam_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']
|
|
# 修改
|
|
sam_img = img
|
|
|
|
if self.use_udp:
|
|
trans = get_warp_matrix(r, c * 2.0, image_size - 1.0, s * 200.0)
|
|
# 修改
|
|
sam_trans = get_warp_matrix(r, c * 2.0, sam_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)
|
|
# 修改
|
|
sam_img = cv2.warpAffine(
|
|
sam_img,
|
|
sam_trans, (int(sam_image_size[0]), int(sam_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
|
|
]
|
|
# 修改
|
|
sam_img = [
|
|
cv2.warpAffine(
|
|
i,
|
|
sam_trans, (int(sam_image_size[0]), int(sam_image_size[1])),
|
|
flags=cv2.INTER_LINEAR) for i in sam_img
|
|
]
|
|
|
|
joints_3d[:, 0:2] = \
|
|
warp_affine_joints(joints_3d[:, 0:2].copy(), trans)
|
|
|
|
else:
|
|
trans = get_affine_transform(c, s, r, image_size)
|
|
# 修改
|
|
sam_trans = get_affine_transform(c, s, r, sam_image_size)
|
|
if not isinstance(img, list):
|
|
img = cv2.warpAffine(
|
|
img,
|
|
trans, (int(image_size[0]), int(image_size[1])),
|
|
flags=cv2.INTER_LINEAR)
|
|
|
|
# 修改
|
|
sam_img = cv2.warpAffine(
|
|
sam_img,
|
|
sam_trans, (int(sam_image_size[0]), int(sam_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
|
|
]
|
|
# 修改
|
|
sam_img = [
|
|
cv2.warpAffine(
|
|
i,
|
|
sam_trans, (int(sam_image_size[0]), int(sam_image_size[1])),
|
|
flags=cv2.INTER_LINEAR) for i in sam_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['sam_img'] = sam_img
|
|
results['joints_3d'] = joints_3d
|
|
results['joints_3d_visible'] = joints_3d_visible
|
|
|
|
return results
|