|
|
|
|
|
|
|
|
|
|
|
|
| from __future__ import absolute_import
|
| from __future__ import division
|
| from __future__ import print_function
|
|
|
| import numpy as np
|
| import cv2
|
| import torch
|
|
|
| class BRG2Tensor_transform(object):
|
| def __call__(self, pic):
|
| img = torch.from_numpy(pic.transpose((2, 0, 1)))
|
| if isinstance(img, torch.ByteTensor):
|
| return img.float()
|
| else:
|
| return img
|
|
|
| class BGR2RGB_transform(object):
|
| def __call__(self, tensor):
|
| return tensor[[2,1,0],:,:]
|
|
|
| def flip_back(output_flipped, matched_parts):
|
| '''
|
| ouput_flipped: numpy.ndarray(batch_size, num_joints, height, width)
|
| '''
|
| assert output_flipped.ndim == 4,\
|
| 'output_flipped should be [batch_size, num_joints, height, width]'
|
|
|
| output_flipped = output_flipped[:, :, :, ::-1]
|
|
|
| for pair in matched_parts:
|
| tmp = output_flipped[:, pair[0], :, :].copy()
|
| output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :]
|
| output_flipped[:, pair[1], :, :] = tmp
|
|
|
| return output_flipped
|
|
|
|
|
| def fliplr_joints(joints, joints_vis, width, matched_parts):
|
| """
|
| flip coords
|
| """
|
|
|
| joints[:, 0] = width - joints[:, 0] - 1
|
|
|
|
|
| for pair in matched_parts:
|
| joints[pair[0], :], joints[pair[1], :] = \
|
| joints[pair[1], :], joints[pair[0], :].copy()
|
| joints_vis[pair[0], :], joints_vis[pair[1], :] = \
|
| joints_vis[pair[1], :], joints_vis[pair[0], :].copy()
|
|
|
| return joints*joints_vis, joints_vis
|
|
|
|
|
| def transform_preds(coords, center, scale, input_size):
|
| target_coords = np.zeros(coords.shape)
|
| trans = get_affine_transform(center, scale, 0, input_size, inv=1)
|
| for p in range(coords.shape[0]):
|
| target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
|
| return target_coords
|
|
|
| def transform_parsing(pred, center, scale, width, height, input_size):
|
|
|
| trans = get_affine_transform(center, scale, 0, input_size, inv=1)
|
| target_pred = cv2.warpAffine(
|
| pred,
|
| trans,
|
| (int(width), int(height)),
|
| flags=cv2.INTER_NEAREST,
|
| borderMode=cv2.BORDER_CONSTANT,
|
| borderValue=(0))
|
|
|
| return target_pred
|
|
|
| def transform_logits(logits, center, scale, width, height, input_size):
|
|
|
| trans = get_affine_transform(center, scale, 0, input_size, inv=1)
|
| channel = logits.shape[2]
|
| target_logits = []
|
| for i in range(channel):
|
| target_logit = cv2.warpAffine(
|
| logits[:,:,i],
|
| trans,
|
| (int(width), int(height)),
|
| flags=cv2.INTER_LINEAR,
|
| borderMode=cv2.BORDER_CONSTANT,
|
| borderValue=(0))
|
| target_logits.append(target_logit)
|
| target_logits = np.stack(target_logits,axis=2)
|
|
|
| return target_logits
|
|
|
|
|
| def get_affine_transform(center,
|
| scale,
|
| rot,
|
| output_size,
|
| shift=np.array([0, 0], dtype=np.float32),
|
| inv=0):
|
| if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
|
| print(scale)
|
| scale = np.array([scale, scale])
|
|
|
| scale_tmp = scale
|
|
|
| src_w = scale_tmp[0]
|
| dst_w = output_size[1]
|
| dst_h = output_size[0]
|
|
|
| rot_rad = np.pi * rot / 180
|
| src_dir = get_dir([0, src_w * -0.5], rot_rad)
|
| dst_dir = np.array([0, (dst_w-1) * -0.5], np.float32)
|
|
|
| src = np.zeros((3, 2), dtype=np.float32)
|
| dst = np.zeros((3, 2), dtype=np.float32)
|
| src[0, :] = center + scale_tmp * shift
|
| src[1, :] = center + src_dir + scale_tmp * shift
|
| dst[0, :] = [(dst_w-1) * 0.5, (dst_h-1) * 0.5]
|
| dst[1, :] = np.array([(dst_w-1) * 0.5, (dst_h-1) * 0.5]) + dst_dir
|
|
|
| src[2:, :] = get_3rd_point(src[0, :], src[1, :])
|
| dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
|
|
|
| if inv:
|
| trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
|
| else:
|
| trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
|
|
|
| return trans
|
|
|
|
|
| def affine_transform(pt, t):
|
| new_pt = np.array([pt[0], pt[1], 1.]).T
|
| new_pt = np.dot(t, new_pt)
|
| return new_pt[:2]
|
|
|
|
|
| def get_3rd_point(a, b):
|
| direct = a - b
|
| return b + np.array([-direct[1], direct[0]], dtype=np.float32)
|
|
|
|
|
| def get_dir(src_point, rot_rad):
|
| sn, cs = np.sin(rot_rad), np.cos(rot_rad)
|
|
|
| src_result = [0, 0]
|
| src_result[0] = src_point[0] * cs - src_point[1] * sn
|
| src_result[1] = src_point[0] * sn + src_point[1] * cs
|
|
|
| return src_result
|
|
|
|
|
| def crop(img, center, scale, output_size, rot=0):
|
| trans = get_affine_transform(center, scale, rot, output_size)
|
|
|
| dst_img = cv2.warpAffine(img,
|
| trans,
|
| (int(output_size[1]), int(output_size[0])),
|
| flags=cv2.INTER_LINEAR)
|
|
|
| return dst_img
|
|
|
