import math import random from collections import Iterable import cv2 import mmcv import numpy as np from detrsmpl.utils.demo_utils import xywh2xyxy, xyxy2xywh from detrsmpl.core.conventions.keypoints_mapping import get_flip_pairs from detrsmpl.utils.transforms import aa_to_rotmat, rotmat_to_aa from ..builder import PIPELINES from .compose import Compose def get_affine_transform(center, scale, rot, output_size, shift=(0., 0.), inv=False, pixel_std=1.0): """Get the affine transform matrix, given the center/scale/rot/output_size. Args: center (np.ndarray[2, ]): Center of the bounding box (x, y). scale (np.ndarray[2, ]): Scale of the bounding box wrt [width, height]. rot (float): Rotation angle (degree). output_size (np.ndarray[2, ] | list(2,)): Size of the destination heatmaps. shift (0-100%): Shift translation ratio wrt the width/height. Default (0., 0.). inv (bool): Option to inverse the affine transform direction. (inv=False: src->dst or inv=True: dst->src) Returns: np.ndarray: The transform matrix. """ assert len(center) == 2 assert len(scale) == 2 assert len(output_size) == 2 assert len(shift) == 2 scale_tmp = scale * pixel_std shift = np.array(shift) src_h = scale_tmp[1] dst_w = output_size[0] dst_h = output_size[1] rot_rad = np.pi * rot / 180 src_dir = rotate_point([0., src_h * -0.5], rot_rad) dst_dir = np.array([0., dst_h * -0.5]) src = np.zeros((3, 2), dtype=np.float32) src[0, :] = center + scale_tmp * shift src[1, :] = center + src_dir + scale_tmp * shift src[2, :] = _get_3rd_point(src[0, :], src[1, :]) dst = np.zeros((3, 2), dtype=np.float32) dst[0, :] = [dst_w * 0.5, dst_h * 0.5] dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir 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, trans_mat): """Apply an affine transformation to the points. Args: pt (np.ndarray): a 2 dimensional point to be transformed trans_mat (np.ndarray): 2x3 matrix of an affine transform Returns: np.ndarray: Transformed points. """ if pt.ndim == 2: new_pt = np.einsum('ij,mj->im', pt, trans_mat) elif pt.ndim == 3: new_pt = np.einsum('nij,mj->nim', pt, trans_mat) else: msg = f'Expected pt to have ndim of 2 or 3, but get {pt.ndim} ' raise ValueError(msg) # new_pt = np.array(trans_mat) @ np.array([pt[0], pt[1], 1.]) return new_pt def _get_3rd_point(a, b): """To calculate the affine matrix, three pairs of points are required. This function is used to get the 3rd point, given 2D points a & b. The 3rd point is defined by rotating vector `a - b` by 90 degrees anticlockwise, using b as the rotation center. Args: a (np.ndarray): point(x,y) b (np.ndarray): point(x,y) Returns: np.ndarray: The 3rd point. """ assert len(a) == 2 assert len(b) == 2 direction = a - b third_pt = b + np.array([-direction[1], direction[0]], dtype=np.float32) return third_pt def rotate_point(pt, angle_rad): """Rotate a point by an angle. Args: pt (list[float]): 2 dimensional point to be rotated angle_rad (float): rotation angle by radian Returns: list[float]: Rotated point. """ assert len(pt) == 2 sn, cs = np.sin(angle_rad), np.cos(angle_rad) new_x = pt[0] * cs - pt[1] * sn new_y = pt[0] * sn + pt[1] * cs rotated_pt = [new_x, new_y] return rotated_pt def get_warp_matrix(theta, size_input, size_dst, size_target): """Calculate the transformation matrix under the constraint of unbiased. Paper ref: Huang et al. The Devil is in the Details: Delving into Unbiased Data Processing for Human Pose Estimation (CVPR 2020). Args: theta (float): Rotation angle in degrees. size_input (np.ndarray): Size of input image [w, h]. size_dst (np.ndarray): Size of output image [w, h]. size_target (np.ndarray): Size of ROI in input plane [w, h]. Returns: matrix (np.ndarray): A matrix for transformation. """ theta = np.deg2rad(theta) matrix = np.zeros((2, 3), dtype=np.float32) scale_x = size_dst[0] / size_target[0] scale_y = size_dst[1] / size_target[1] matrix[0, 0] = math.cos(theta) * scale_x matrix[0, 1] = -math.sin(theta) * scale_x matrix[0, 2] = scale_x * (-0.5 * size_input[0] * math.cos(theta) + 0.5 * size_input[1] * math.sin(theta) + 0.5 * size_target[0]) matrix[1, 0] = math.sin(theta) * scale_y matrix[1, 1] = math.cos(theta) * scale_y matrix[1, 2] = scale_y * (-0.5 * size_input[0] * math.sin(theta) - 0.5 * size_input[1] * math.cos(theta) + 0.5 * size_target[1]) return matrix def warp_affine_joints(joints, mat): """Apply affine transformation defined by the transform matrix on the joints. Args: joints (np.ndarray[..., 2]): Origin coordinate of joints. mat (np.ndarray[3, 2]): The affine matrix. Returns: matrix (np.ndarray[..., 2]): Result coordinate of joints. """ joints = np.array(joints) shape = joints.shape joints = joints.reshape(-1, 2) return np.dot(np.concatenate((joints, joints[:, 0:1] * 0 + 1), axis=1), mat.T).reshape(shape) 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 _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_smplx_pose(pose): """Flip SMPLX pose parameters horizontally. Args: pose (np.ndarray([63])): SMPLX pose parameters Returns: pose_flipped (np.ndarray([21,3])) """ flippedParts = np.array([ 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 ], dtype=np.int32) - 3 dim_flip = np.array([1, -1, -1], dtype=pose.dtype) pose = (pose[..., flippedParts].reshape(-1, 21, 3) * dim_flip).copy() return pose def _flip_axis_angle(r): """Flip axis_angle horizontally. Args: r (np.ndarray([3])) Returns: f_flipped """ dim_flip = np.array([1, -1, -1], dtype=r.dtype) r = r * dim_flip return r def _flip_hand_pose(r_pose, l_pose): dim_flip = np.array([1, -1, -1], dtype=r_pose.dtype) ret_l_pose = r_pose * dim_flip ret_r_pose = l_pose * dim_flip return ret_r_pose, ret_l_pose def _flip_keypoints(keypoints, flip_pairs, img_width=None): """Flip human joints horizontally. Note: num_keypoints: K num_dimension: D Args: keypoints (np.ndarray([K, D])): Coordinates of keypoints. flip_pairs (list[tuple()]): Pairs of keypoints which are mirrored (for example, left ear -- right ear). img_width (int | None, optional): The width of the original image. To flip 2D keypoints, image width is needed. To flip 3D keypoints, we simply negate the value of x-axis. Default: None. Returns: keypoints_flipped """ keypoints_flipped = keypoints.copy() # Swap left-right parts for left, right in flip_pairs: keypoints_flipped[..., left, :] = keypoints[..., right, :] keypoints_flipped[..., right, :] = keypoints[..., left, :] # Flip horizontally if img_width is None: keypoints_flipped[..., 0] = -keypoints_flipped[..., 0] else: keypoints_flipped[..., 0] = img_width - 1 - keypoints_flipped[..., 0] return keypoints_flipped def _rotate_joints_3d(joints_3d, rot): """Rotate the 3D joints in the local coordinates. Notes: 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) if joints_3d.ndim == 2: joints_3d_rotated = np.einsum('ij,kj->ki', rot_mat, joints_3d) elif joints_3d.ndim == 3: joints_3d_rotated = np.einsum('ij,mkj->mki', rot_mat, joints_3d) else: msg = 'Expected joints_3d to have ndim of 2 or 3, ' f'but get {joints_3d.ndim}.' raise ValueError(msg) 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.astype(np.float32)) # # 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] # use pytorch3d rot_mat = _construct_rotation_matrix(-rot) orient = pose[..., :3] per_rdg = aa_to_rotmat(orient) if pose.ndim == 1: tmp_rot = np.einsum('ij,jk->ik', rot_mat, per_rdg) elif pose.ndim == 2: tmp_rot = np.einsum('ij,mjk->mik', rot_mat, per_rdg) else: msg = f'Expected pose to have ndim of 2 or 3, but get {pose.ndim} ' raise ValueError(msg) res_rot = rotmat_to_aa(tmp_rot) pose_rotated[..., :3] = res_rot # use cv2 # rot_mat = _construct_rotation_matrix(-rot) # for i in range(pose.shape[0]): # orient = pose[i, :3] # # find the rotation of the body in camera frame # per_rdg, _ = cv2.Rodrigues(orient.astype(np.float32)) # # apply the global rotation to the global orientation # res_rot, _ = cv2.Rodrigues(np.dot(rot_mat, per_rdg)) # pose_rotated[i, :3] = (res_rot.T)[0] return pose_rotated def _bbox_flip(bboxes, img_shape, direction): """Flip bboxes horizontally. Args: bboxes (numpy.ndarray): Bounding boxes, shape (..., 4*k) img_shape (tuple[int]): Image shape (height, width) direction (str): Flip direction. Options are 'horizontal', 'vertical'. Returns: numpy.ndarray: Flipped bounding boxes. """ assert bboxes.shape[-1] % 5 == 0 flipped = bboxes.copy() if direction == 'horizontal': w = img_shape[1] flipped[..., 0::4] = w - bboxes[..., 2::4] flipped[..., 2::4] = w - bboxes[..., 0::4] elif direction == 'vertical': h = img_shape[0] flipped[..., 1::4] = h - bboxes[..., 3::4] flipped[..., 3::4] = h - bboxes[..., 1::4] elif direction == 'diagonal': w = img_shape[1] h = img_shape[0] flipped[..., 0::4] = w - bboxes[..., 2::4] flipped[..., 1::4] = h - bboxes[..., 3::4] flipped[..., 2::4] = w - bboxes[..., 0::4] flipped[..., 3::4] = h - bboxes[..., 1::4] else: raise ValueError(f"Invalid flipping direction '{direction}'") return flipped @PIPELINES.register_module() class RandomHorizontalFlip(object): """Flip the image randomly. Flip the image randomly based on flip probaility. Args: flip_prob (float): probability of the image being flipped. Default: 0.5 """ def __init__(self, flip_prob=0.5, convention=None): assert 0 <= flip_prob <= 1 self.flip_prob = flip_prob self.flip_pairs = get_flip_pairs(convention) def __call__(self, results): """Call function to flip image and annotations. Args: results (dict): Result dict from loading pipeline. Returns: dict: Flipped results, 'flip' key is added into result dict. """ if np.random.rand() > self.flip_prob: results['is_flipped'] = np.array([0]) return results results['is_flipped'] = np.array([1]) # flip image for key in results.get('img_fields', ['img']): results[key] = mmcv.imflip(results[key], direction='horizontal') # flip keypoints2d if 'keypoints2d' in results: assert self.flip_pairs is not None width = results['img'][:, ::-1, :].shape[1] keypoints2d = results['keypoints2d'].copy() keypoints2d = _flip_keypoints(keypoints2d, self.flip_pairs, width) results['keypoints2d'] = keypoints2d elif 'keypoints2d_ori' in results: assert self.flip_pairs is not None width = results['img'][:, ::-1, :].shape[1] keypoints2d = results['keypoints2d_ori'].copy() keypoints2d = _flip_keypoints(keypoints2d, self.flip_pairs, width) results['keypoints2d_ori'] = keypoints2d if 'keypoints2d_smpl' in results: assert self.flip_pairs is not None width = results['img'][:, ::-1, :].shape[1] keypoints2d = results['keypoints2d_smpl'].copy() keypoints2d = _flip_keypoints(keypoints2d, self.flip_pairs, width) results['keypoints2d_smpl'] = keypoints2d # flip bbox center center = results['center'] center[..., 0] = width - 1 - center[..., 0] results['center'] = center # flip keypoints3d if 'keypoints3d' in results: assert self.flip_pairs is not None keypoints3d = results['keypoints3d'].copy() keypoints3d = _flip_keypoints(keypoints3d, self.flip_pairs) results['keypoints3d'] = keypoints3d elif 'keypoints3d_ori' in results: assert self.flip_pairs is not None keypoints3d = results['keypoints3d_ori'].copy() keypoints3d = _flip_keypoints(keypoints3d, self.flip_pairs) results['keypoints3d_ori'] = keypoints3d if 'keypoints3d_smpl' in results: assert self.flip_pairs is not None keypoints3d = results['keypoints3d_smpl'].copy() keypoints3d = _flip_keypoints(keypoints3d, self.flip_pairs) results['keypoints3d_smpl'] = keypoints3d if 'bbox_xywh' in results: width = results['img'].shape[1] bbox_xywh = results['bbox_xywh'].copy() bbox_xyxy = xywh2xyxy(bbox_xywh) bbox_xyxy = bbox_xyxy[:, [2, 1, 0, 3]] * np.array( [-1, 1, -1, 1]) + np.array([width, 0, width, 0]) # img = mmcv.imshow_bboxes(results['img'], bbox_xyxy, show=False) # cv2.imwrite('test.png',img) results['bbox_xywh'] = xyxy2xywh(bbox_xyxy) # flip smpl if 'smpl_body_pose' in results: global_orient = results['smpl_global_orient'].copy() body_pose = results['smpl_body_pose'].copy().reshape((-1, 23 * 3)) smpl_pose = np.concatenate((global_orient, body_pose), axis=-1) smpl_pose_flipped = _flip_smpl_pose(smpl_pose) global_orient = smpl_pose_flipped[..., :3] body_pose = smpl_pose_flipped[..., 3:] results['smpl_global_orient'] = global_orient results['smpl_body_pose'] = body_pose.reshape((-1, 23, 3)) # TODO: to check multi-human for smplx if 'smplx_body_pose' in results: body_pose = results['smplx_body_pose'].copy().reshape((-1)) body_pose_flipped = _flip_smplx_pose(body_pose) results['smplx_body_pose'] = body_pose_flipped if 'smplx_global_orient' in results: global_orient = results['smplx_global_orient'].copy().reshape((-1)) global_orient_flipped = _flip_axis_angle(global_orient) results['smplx_global_orient'] = global_orient_flipped if 'smplx_jaw_pose' in results: jaw_pose = results['smplx_jaw_pose'].copy().reshape((-1)) jaw_pose_flipped = _flip_axis_angle(jaw_pose) results['smplx_jaw_pose'] = jaw_pose_flipped if 'smplx_right_hand_pose' in results: right_hand_pose = results['smplx_right_hand_pose'].copy() left_hand_pose = results['smplx_left_hand_pose'].copy() results['smplx_right_hand_pose'], results[ 'smplx_left_hand_pose'] = _flip_hand_pose( right_hand_pose, left_hand_pose) # Expressions are not symmetric. Remove them when flipped. if 'smplx_expression' in results: results['smplx_expression'] = np.zeros( (results['smplx_expression'].shape[0]), dtype=np.float32) results['has_smplx_expression'] = 0 return results def __repr__(self): return self.__class__.__name__ + f'(flip_prob={self.flip_prob})' def resize(ori_shape, size, max_size=None): # size can be min_size (scalar) or (w, h) tuple # import ipdb; ipdb.set_trace(context=15) def get_size_with_aspect_ratio(image_size, size, max_size=None): w, h = image_size if max_size is not None: min_original_size = float(min((w, h))) max_original_size = float(max((w, h))) if max_original_size / min_original_size * size > max_size: size = int( round(max_size * min_original_size / max_original_size)) if (w <= h and w == size) or (h <= w and h == size): return (w, h) if w < h: ow = size oh = int(size * h / w) else: oh = size ow = int(size * w / h) return (ow, oh) def get_size(ori_shape, size, max_size=None): if isinstance(size, (list, tuple)): return size[::-1] else: return get_size_with_aspect_ratio(ori_shape, size, max_size) size = get_size(ori_shape, size, max_size) return size @PIPELINES.register_module() class CenterCrop(object): r"""Center crop the image. Args: crop_size (int | tuple): Expected size after cropping with the format of (h, w). efficientnet_style (bool): Whether to use efficientnet style center crop. Defaults to False. crop_padding (int): The crop padding parameter in efficientnet style center crop. Only valid if efficientnet style is True. Defaults to 32. interpolation (str): Interpolation method, accepted values are 'nearest', 'bilinear', 'bicubic', 'area', 'lanczos'. Only valid if efficientnet style is True. Defaults to 'bilinear'. backend (str): The image resize backend type, accpeted values are `cv2` and `pillow`. Only valid if efficientnet style is True. Defaults to `cv2`. Notes: If the image is smaller than the crop size, return the original image. If efficientnet_style is set to False, the pipeline would be a simple center crop using the crop_size. If efficientnet_style is set to True, the pipeline will be to first to perform the center crop with the crop_size_ as: .. math:: crop\_size\_ = crop\_size / (crop\_size + crop\_padding) * short\_edge And then the pipeline resizes the img to the input crop size. """ def __init__(self, crop_size, efficientnet_style=False, crop_padding=32, interpolation='bilinear', backend='cv2'): if efficientnet_style: assert isinstance(crop_size, int) assert crop_padding >= 0 assert interpolation in ('nearest', 'bilinear', 'bicubic', 'area', 'lanczos') if backend not in ['cv2', 'pillow']: raise ValueError( f'backend: {backend} is not supported for ' 'resize. Supported backends are "cv2", "pillow"') else: assert isinstance(crop_size, int) or (isinstance(crop_size, tuple) and len(crop_size) == 2) if isinstance(crop_size, int): crop_size = (crop_size, crop_size) assert crop_size[0] > 0 and crop_size[1] > 0 self.crop_size = crop_size self.efficientnet_style = efficientnet_style self.crop_padding = crop_padding self.interpolation = interpolation self.backend = backend def __call__(self, results): crop_height, crop_width = self.crop_size[0], self.crop_size[1] for key in results.get('img_fields', ['img']): img = results[key] # img.shape has length 2 for grayscale, length 3 for color img_height, img_width = img.shape[:2] # https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/preprocessing.py#L118 # noqa if self.efficientnet_style: img_short = min(img_height, img_width) crop_height = crop_height / (crop_height + self.crop_padding) * img_short crop_width = crop_width / (crop_width + self.crop_padding) * img_short y1 = max(0, int(round((img_height - crop_height) / 2.))) x1 = max(0, int(round((img_width - crop_width) / 2.))) y2 = min(img_height, y1 + crop_height) - 1 x2 = min(img_width, x1 + crop_width) - 1 # crop the image img = mmcv.imcrop(img, bboxes=np.array([x1, y1, x2, y2])) if self.efficientnet_style: img = mmcv.imresize(img, tuple(self.crop_size[::-1]), interpolation=self.interpolation, backend=self.backend) img_shape = img.shape results[key] = img results['img_shape'] = img_shape return results def __repr__(self): repr_str = self.__class__.__name__ + f'(crop_size={self.crop_size}' repr_str += f', efficientnet_style={self.efficientnet_style}' repr_str += f', crop_padding={self.crop_padding}' repr_str += f', interpolation={self.interpolation}' repr_str += f', backend={self.backend})' return repr_str @PIPELINES.register_module() class Normalize(object): """Normalize the image. Args: mean (sequence): Mean values of 3 channels. std (sequence): Std values of 3 channels. to_rgb (bool): Whether to convert the image from BGR to RGB, default is true. """ def __init__(self, mean, std, to_rgb=True): self.mean = np.array(mean, dtype=np.float32) self.std = np.array(std, dtype=np.float32) self.to_rgb = to_rgb def __call__(self, results): for key in results.get('img_fields', ['img']): results[key] = mmcv.imnormalize(results[key], self.mean, self.std, self.to_rgb) results['img_norm_cfg'] = dict(mean=self.mean, std=self.std, to_rgb=self.to_rgb) return results def __repr__(self): repr_str = self.__class__.__name__ repr_str += f'(mean={list(self.mean)}, ' repr_str += f'std={list(self.std)}, ' repr_str += f'to_rgb={self.to_rgb})' return repr_str @PIPELINES.register_module() class ColorJitter(object): """Randomly change the brightness, contrast and saturation of an image. Args: brightness (float): How much to jitter brightness. brightness_factor is chosen uniformly from [max(0, 1 - brightness), 1 + brightness]. contrast (float): How much to jitter contrast. contrast_factor is chosen uniformly from [max(0, 1 - contrast), 1 + contrast]. saturation (float): How much to jitter saturation. saturation_factor is chosen uniformly from [max(0, 1 - saturation), 1 + saturation]. """ def __init__(self, brightness, contrast, saturation): self.brightness = brightness self.contrast = contrast self.saturation = saturation def __call__(self, results): brightness_factor = random.uniform(0, self.brightness) contrast_factor = random.uniform(0, self.contrast) saturation_factor = random.uniform(0, self.saturation) color_jitter_transforms = [ dict(type='Brightness', magnitude=brightness_factor, prob=1., random_negative_prob=0.5), dict(type='Contrast', magnitude=contrast_factor, prob=1., random_negative_prob=0.5), dict(type='ColorTransform', magnitude=saturation_factor, prob=1., random_negative_prob=0.5) ] random.shuffle(color_jitter_transforms) transform = Compose(color_jitter_transforms) return transform(results) def __repr__(self): repr_str = self.__class__.__name__ repr_str += f'(brightness={self.brightness}, ' repr_str += f'contrast={self.contrast}, ' repr_str += f'saturation={self.saturation})' return repr_str @PIPELINES.register_module() class Lighting(object): """Adjust images lighting using AlexNet-style PCA jitter. Args: eigval (list): the eigenvalue of the convariance matrix of pixel values, respectively. eigvec (list[list]): the eigenvector of the convariance matrix of pixel values, respectively. alphastd (float): The standard deviation for distribution of alpha. Defaults to 0.1 to_rgb (bool): Whether to convert img to rgb. """ def __init__(self, eigval, eigvec, alphastd=0.1, to_rgb=True): assert isinstance(eigval, list), \ f'eigval must be of type list, got {type(eigval)} instead.' assert isinstance(eigvec, list), \ f'eigvec must be of type list, got {type(eigvec)} instead.' for vec in eigvec: assert isinstance(vec, list) and len(vec) == len(eigvec[0]), \ 'eigvec must contains lists with equal length.' self.eigval = np.array(eigval) self.eigvec = np.array(eigvec) self.alphastd = alphastd self.to_rgb = to_rgb def __call__(self, results): for key in results.get('img_fields', ['img']): img = results[key] results[key] = mmcv.adjust_lighting(img, self.eigval, self.eigvec, alphastd=self.alphastd, to_rgb=self.to_rgb) return results def __repr__(self): repr_str = self.__class__.__name__ repr_str += f'(eigval={self.eigval.tolist()}, ' repr_str += f'eigvec={self.eigvec.tolist()}, ' repr_str += f'alphastd={self.alphastd}, ' repr_str += f'to_rgb={self.to_rgb})' return repr_str @PIPELINES.register_module() class RandomChannelNoise: """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 if 'ori_img' in results: img = results['ori_img'] img = cv2.multiply(img, pn) results['ori_img'] = img return results @PIPELINES.register_module() class GetRandomScaleRotation: """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.0 results['scale'] = s results['rotation'] = r return results @PIPELINES.register_module() class SampleInstance: def __init__(self, sample_ratio): self.sample_ratio = sample_ratio def __call__(self, results): assert 'bbox_xywh' in results bbox_xywh = results['bbox_xywh'].copy() crop_person_number = len(bbox_xywh) if random.random() < self.sample_ratio: crop_person_number = np.random.randint(len(bbox_xywh)) + 1 sample_ids = np.array( random.sample(list(range(len(bbox_xywh))), crop_person_number)) bbox_xyxy = xywh2xyxy(bbox_xywh)[sample_ids] leftTop_ = bbox_xyxy[:, :2] leftTop_ = np.array([np.min(leftTop_[:, 0]), np.min(leftTop_[:, 1])]) rightBottom_ = bbox_xyxy[:, 2:4] rightBottom_ = np.array( [np.max(rightBottom_[:, 0]), np.max(rightBottom_[:, 1])]) bbox_xyxy = np.concatenate([leftTop_, rightBottom_]) results['bbox_xyxy'] = bbox_xyxy center = (rightBottom_ + leftTop_) / 2 scale = (rightBottom_ - leftTop_) scale[0] = scale[1] = max(scale) results['center'] = center results['scale'] = scale return results @PIPELINES.register_module() class MeshAffine: """Affine transform the image to get input image. Affine transform the 2D keypoints, 3D kepoints. Required keys: 'img', 'pose', 'img_shape', 'rotation' and 'center'. Modifies key: 'img', ''keypoints2d', 'keypoints3d', 'pose'. """ def __init__(self, img_res, crop_with_bbox=True): if isinstance(img_res, tuple): self.image_size = img_res else: self.image_size = np.array([img_res, img_res]) self.img_res = img_res self.crop_with_bbox = crop_with_bbox def __call__(self, results): c = results['center'] s = results['scale'] r = results['rotation'] trans = get_affine_transform(c, s, r, self.image_size) if 'img' in results: img = results['img'].copy() # img before affine ori_img = img.copy() results['crop_transform'] = trans results['ori_img'] = ori_img results['img_fields'] = ['img', 'ori_img'] img = cv2.warpAffine( img, trans, (int(self.image_size[0]), int(self.image_size[1])), flags=cv2.INTER_LINEAR) results['img'] = img if 'keypoints2d' in results: keypoints2d = results['keypoints2d'].copy() results['keypoints2d'][..., :2] = affine_transform( keypoints2d, trans) if 'bbox_xywh' in results: bbox_xywh = results['bbox_xywh'].copy() leftTop = bbox_xywh[..., :2] rightTop = np.concatenate([ bbox_xywh[..., [0]] + bbox_xywh[..., [2]], bbox_xywh[..., [1]] ], -1) leftBottom = np.concatenate([ bbox_xywh[..., [0]], bbox_xywh[..., [1]] + bbox_xywh[..., [3]] ], -1) rightBottom = np.concatenate([ bbox_xywh[..., [0]] + bbox_xywh[..., [2]], bbox_xywh[..., [1]] + bbox_xywh[..., [3]] ], -1) bbox_point = np.vstack( [leftTop, rightTop, leftBottom, rightBottom]) bbox_point = np.concatenate( [bbox_point, np.ones_like(bbox_point[..., [0]])], -1) bbox_point = affine_transform(bbox_point, trans) # TODO: bbox_point = np.clip(bbox_point, 0, self.img_res) bbox__xywh_t = bbox_point.clone() bbox__xywh_t results['bbox'] = bbox_point # bbox_xyxy = xywh2xyxy(bbox_xywh)[:,:4].reshape(-1, 2, 2) # bbox_xyxy = np.concatenate([bbox_xyxy, np.ones_like(bbox_xyxy[...,[0]])], -1) # bbox_xyxy = np.concatenate([affine_transform(bbox_xyxy, trans).reshape(-1,4), bbox_xywh[...,[-1]]],-1) # results['bbox_xywh'] = xyxy2xywh(bbox_xyxy) # image_array=np.array([img]), # overwrite=True, # data_source='smpl_54') if 'keypoints3d' in results: keypoints3d = results['keypoints3d'].copy() keypoints3d[..., :3] = _rotate_joints_3d(keypoints3d[..., :3], r) results['keypoints3d'] = keypoints3d if 'smpl_body_pose' in results: global_orient = results['smpl_global_orient'].copy() body_pose = results['smpl_body_pose'].copy().reshape((-1, 23 * 3)) pose = np.concatenate((global_orient, body_pose), axis=-1) pose = _rotate_smpl_pose(pose, r) results['smpl_global_orient'] = pose[..., :3] results['smpl_body_pose'] = pose[..., 3:].reshape((-1, 23, 3)) if 'smplx_global_orient' in results: global_orient = results['smplx_global_orient'].copy() global_orient = _rotate_smpl_pose(global_orient, r) results['smplx_global_orient'] = global_orient return results @PIPELINES.register_module() class MeshAffineED: """Affine transform the image to get input image. Affine transform the 2D keypoints, 3D kepoints. Required keys: 'img', 'pose', 'img_shape', 'rotation' and 'center'. Modifies key: 'img', ''keypoints2d', 'keypoints3d', 'pose'. """ def __init__(self, sizes, max_size=None): assert isinstance(sizes, (list, tuple)) self.sizes = sizes self.max_size = max_size def __call__(self, results): ori_shape = np.array(results['ori_shape']) # ori_shape = ori_shape[::-1] # print(ori_shape) size = random.choice(self.sizes) reshape_size = resize(ori_shape, size, self.max_size) c = (ori_shape / 2)[::-1] s = ori_shape[::-1] r = results['rotation'] trans = get_affine_transform(c, s, r, reshape_size[::-1]) results['img_shape'] = reshape_size if 'img' in results: img = results['img'].copy() # img before affine ori_img = img.copy() results['crop_transform'] = trans results['ori_img'] = ori_img results['img_fields'] = ['img', 'ori_img'] img = cv2.warpAffine(img, trans, (int(reshape_size[1]), int(reshape_size[0])), flags=cv2.INTER_LINEAR) results['img'] = img if 'keypoints2d_ori' in results: keypoints2d_ori = results['keypoints2d_ori'].copy() results['keypoints2d_ori'][..., :2] = affine_transform( keypoints2d_ori, trans) if 'keypoints2d_smpl' in results: keypoints2d_smpl = results['keypoints2d_smpl'].copy() results['keypoints2d_smpl'][..., :2] = affine_transform( keypoints2d_smpl, trans) if 'bbox_xywh' in results: bbox_xywh = results['bbox_xywh'].copy() leftTop = bbox_xywh[..., :2] rightTop = np.concatenate([ bbox_xywh[..., [0]] + bbox_xywh[..., [2]], bbox_xywh[..., [1]] ], -1) leftBottom = np.concatenate([ bbox_xywh[..., [0]], bbox_xywh[..., [1]] + bbox_xywh[..., [3]] ], -1) rightBottom = np.concatenate([ bbox_xywh[..., [0]] + bbox_xywh[..., [2]], bbox_xywh[..., [1]] + bbox_xywh[..., [3]] ], -1) bbox_point = np.vstack( [leftTop, rightTop, leftBottom, rightBottom]) bbox_point = np.concatenate( [bbox_point, np.ones_like(bbox_point[..., [0]])], -1) bbox_point = affine_transform(bbox_point, trans) # TODO: bbox_point = np.clip(bbox_point, 0, (int(reshape_size[1]), int(reshape_size[0]))) results['bbox'] = bbox_point bbox_xyxy_t = bbox_xywh.copy() num_sample = bbox_xywh.shape[0] bbox_xyxy_t[..., :2] = bbox_point[:num_sample, :] bbox_xyxy_t[..., 2:4] = bbox_point[num_sample * 3:num_sample * 4, :] results['bbox_xywh'] = xyxy2xywh(bbox_xyxy_t) # bbox_xywh = results['bbox_xywh'].copy() # bbox_xyxy = xywh2xyxy(bbox_xywh)[:,:4].reshape(-1, 2, 2) # bbox_xyxy = np.concatenate([bbox_xyxy, np.ones_like(bbox_xyxy[...,[0]])], -1) # bbox_xyxy = np.concatenate([affine_transform(bbox_xyxy, trans).reshape(-1,4), bbox_xywh[...,[-1]]],-1) # results['bbox_xywh'] = xyxy2xywh(bbox_xyxy) # image_array=np.array([img]), # overwrite=True, # data_source='smpl_54') if 'keypoints3d_ori' in results: keypoints3d_ori = results['keypoints3d_ori'].copy() keypoints3d_ori[..., :3] = _rotate_joints_3d( keypoints3d_ori[..., :3], r) results['keypoints3d_ori'] = keypoints3d_ori if 'keypoints3d_smpl' in results: keypoints3d_smpl = results['keypoints3d_smpl'].copy() keypoints3d_smpl[..., :3] = _rotate_joints_3d( keypoints3d_smpl[..., :3], r) results['keypoints3d_smpl'] = keypoints3d_smpl if 'smpl_body_pose' in results: global_orient = results['smpl_global_orient'].copy() body_pose = results['smpl_body_pose'].copy().reshape((-1, 23 * 3)) pose = np.concatenate((global_orient, body_pose), axis=-1) pose = _rotate_smpl_pose(pose, r) results['smpl_global_orient'] = pose[..., :3] results['smpl_body_pose'] = pose[..., 3:].reshape((-1, 23, 3)) if 'area' in results: area = results['area'] * (trans[0, 0] * trans[1, 1]) results['area'] = area # if 'smplx_global_orient' in results: # global_orient = results['smplx_global_orient'].copy() # global_orient = _rotate_smpl_pose(global_orient, r) # results['smplx_global_orient'] = global_orient return results @PIPELINES.register_module() class Rotation: """Rotate the image with the given rotation. Rotate the 2D keypoints, 3D kepoints, poses. Required keys: 'img', 'pose', 'rotation' and 'center'. Modifies key: 'img', ''keypoints2d', 'keypoints3d', 'pose'. To avoid conflicts with MeshAffine, rotation will be set to 0.0 after rotate the image. The rotation value will be stored to 'ori_rotation'. """ def __init__(self): pass def __call__(self, results): r = results['rotation'] if r == 0.0: return results img = results['img'] # img before affine (h, w) = img.shape[:2] (cX, cY) = (w // 2, h // 2) M = cv2.getRotationMatrix2D((cX, cY), r, 1.0) cos = np.abs(M[0, 0]) sin = np.abs(M[0, 1]) # compute the new bounding dimensions of the image nW = int((h * sin) + (w * cos)) nH = int((h * cos) + (w * sin)) # adjust the rotation matrix to take into account translation M[0, 2] += (nW / 2) - cX M[1, 2] += (nH / 2) - cY # perform the actual rotation and return the image img = cv2.warpAffine(img, M, (nW, nH)) results['img'] = img c = results['center'] c = np.dot(M[:2, :2], c) + M[:2, 2] results['center'] = c if 'keypoints2d' in results: keypoints2d = results['keypoints2d'].copy() keypoints2d[:, :2] = (np.dot(keypoints2d[:, :2], M[:2, :2].T) + M[:2, 2] + 1).astype(np.int) results['keypoints2d'] = keypoints2d if 'keypoints3d' in results: keypoints3d = results['keypoints3d'].copy() keypoints3d[:, :3] = _rotate_joints_3d(keypoints3d[:, :3], r) results['keypoints3d'] = keypoints3d if 'smpl_body_pose' in results: global_orient = results['smpl_global_orient'].copy() body_pose = results['smpl_body_pose'].copy().reshape((-1)) pose = np.concatenate((global_orient, body_pose), axis=-1) pose = _rotate_smpl_pose(pose, r) results['smpl_global_orient'] = pose[:3] results['smpl_body_pose'] = pose[3:].reshape((-1, 3)) if 'smplx_global_orient' in results: global_orient = results['smplx_global_orient'].copy() global_orient = _rotate_smpl_pose(global_orient, r) results['smplx_global_orient'] = global_orient results['rotation'] = 0.0 results['ori_rotation'] = r return results @PIPELINES.register_module() class BBoxCenterJitter(object): def __init__(self, factor=0.0, dist='normal'): super(BBoxCenterJitter, self).__init__() self.factor = factor self.dist = dist assert self.dist in [ 'normal', 'uniform' ], (f'Distribution must be normal or uniform, not {self.dist}') def __call__(self, results): # body model: no process if self.factor <= 1e-3: return results bbox_size = results['scale'][0] jitter = bbox_size * self.factor if self.dist == 'normal': center_jitter = np.random.randn(2) * jitter elif self.dist == 'uniform': center_jitter = np.random.rand(2) * 2 * jitter - jitter center = results['center'] H, W = results['img_shape'] new_center = center + center_jitter new_center[0] = np.clip(new_center[0], 0, W) new_center[1] = np.clip(new_center[1], 0, H) results['center'] = new_center return results @PIPELINES.register_module() class SimulateLowRes(object): def __init__(self, dist: str = 'categorical', factor: float = 1.0, cat_factors=(1.0, ), factor_min: float = 1.0, factor_max: float = 1.0) -> None: self.factor_min = factor_min self.factor_max = factor_max self.dist = dist self.cat_factors = cat_factors assert dist in ['uniform', 'categorical'] def _sample_low_res(self, image: np.ndarray) -> np.ndarray: """""" if self.dist == 'uniform': downsample = self.factor_min != self.factor_max if not downsample: return image factor = np.random.rand() * (self.factor_max - self.factor_min) + self.factor_min elif self.dist == 'categorical': if len(self.cat_factors) < 2: return image idx = np.random.randint(0, len(self.cat_factors)) factor = self.cat_factors[idx] H, W, _ = image.shape downsampled_image = cv2.resize(image, (int(W // factor), int(H // factor)), cv2.INTER_NEAREST) resized_image = cv2.resize(downsampled_image, (W, H), cv2.INTER_LINEAR_EXACT) return resized_image def __call__(self, results): """""" img = results['img'] img = self._sample_low_res(img) results['img'] = img return results