import numpy as np import cv2 import random from config.config import cfg import math from .human_models import smpl_x, smpl from .transforms import cam2pixel, transform_joint_to_other_db, transform_joint_to_other_db_batch from plyfile import PlyData, PlyElement import torch import torch.distributed as dist def load_img(path, order='RGB'): img = cv2.imread(path, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION) if not isinstance(img, np.ndarray): raise IOError('Fail to read %s' % path) if order == 'RGB': img = img[:, :, ::-1].copy() img = img.astype(np.float32) return img def get_bbox(joint_img, joint_valid, extend_ratio=1.2): x_img, y_img = joint_img[:, 0], joint_img[:, 1] x_img = x_img[joint_valid == 1] y_img = y_img[joint_valid == 1] xmin = min(x_img) ymin = min(y_img) xmax = max(x_img) ymax = max(y_img) x_center = (xmin + xmax) / 2. width = xmax - xmin xmin = x_center - 0.5 * width * extend_ratio xmax = x_center + 0.5 * width * extend_ratio y_center = (ymin + ymax) / 2. height = ymax - ymin ymin = y_center - 0.5 * height * extend_ratio ymax = y_center + 0.5 * height * extend_ratio bbox = np.array([xmin, ymin, xmax - xmin, ymax - ymin]).astype(np.float32) return bbox def sanitize_bbox(bbox, img_width, img_height): x, y, w, h = bbox x1 = np.max((0, x)) y1 = np.max((0, y)) x2 = np.min((img_width - 1, x1 + np.max((0, w - 1)))) y2 = np.min((img_height - 1, y1 + np.max((0, h - 1)))) if w > 0 and h > 0 and x2 >= x1 and y2 >= y1: bbox = np.array([x1, y1, x2 - x1 + 1, y2 - y1 + 1]) else: bbox = None return bbox 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 min_original_size ==0: print('min_original_size:',min_original_size) 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 def process_bbox(bbox, img_width, img_height, ratio=1.): bbox = np.array(bbox, dtype=np.float32) # aspect ratio preserving bbox w = bbox[2] h = bbox[3] c_x = bbox[0] + w / 2. c_y = bbox[1] + h / 2. bbox[2] = w * ratio bbox[3] = h * ratio bbox[0] = c_x - bbox[2] / 2. bbox[1] = c_y - bbox[3] / 2. bbox = sanitize_bbox(bbox, img_width, img_height) return bbox def get_aug_config(data_name): scale_factor = 0.25 rot_factor = 30 color_factor = 0.2 crop_factor = 0.1 if data_name == 'GTA_Human2': sample_ratio = 0.5 sample_prob = 0.5 elif data_name == 'AGORA_MM': sample_ratio = 0.5 sample_prob = 0.7 elif data_name == 'BEDLAM': sample_ratio = 0.6 sample_prob = 0.7 elif data_name == 'Synbody': sample_ratio = 0.6 sample_prob = 0.7 elif data_name == 'COCO_NA': sample_ratio = 0.6 sample_prob = 0.7 elif data_name == 'CrowdPose': sample_ratio = 0.5 sample_prob = 0.5 elif data_name == 'PoseTrack': sample_ratio = 0.5 sample_prob = 0.3 elif data_name == 'UBody_MM': sample_ratio = 0.5 sample_prob = 0.3 elif data_name == 'ARCTIC': sample_ratio = 0.5 sample_prob = 0.3 elif data_name == 'RICH': sample_ratio = 0.5 sample_prob = 0.3 elif data_name == 'EgoBody_Egocentric': sample_ratio = 0.5 sample_prob = 0.3 elif data_name == 'EgoBody_Kinect': sample_ratio = 0.5 sample_prob = 0.3 else: sample_ratio = 0.5 sample_prob = 0.3 scale = np.clip(np.random.randn(), -1.0, 1.0) * scale_factor + 1.0 rot = np.clip(np.random.randn(), -2.0, 2.0) * rot_factor if random.random() <= 0.6 else 0 c_up = 1.0 + color_factor c_low = 1.0 - color_factor color_scale = np.array([ random.uniform(c_low, c_up), random.uniform(c_low, c_up), random.uniform(c_low, c_up) ]) do_flip = random.random() < 0.5 crop_hw = np.array([ 0.2 - np.random.rand() * crop_factor, 0.2 - np.random.rand() * crop_factor ]) # crop_hw = np.array([ # 0.3 - np.random.rand() * crop_factor, 0.3 - np.random.rand() * crop_factor # ]) return scale, rot, color_scale, do_flip, crop_hw, sample_ratio, sample_prob def augmentation_keep_size(img, bbox, data_split): ori_shape = img.shape[:2][::-1] if getattr(cfg, 'no_aug', False) and data_split == 'train': scale, rot, color_scale, do_flip,size,crop = 1.0, 0.0, np.array([1, 1, 1]), False, ori_shape, np.array([1,1]) size = random.choice(cfg.train_sizes) max_size = cfg.train_max_size elif data_split == 'train': scale, rot, color_scale, do_flip, crop = get_aug_config() rot=0 # scale, rot, do_flip, crop = 1.0, 0.0, False, np.array([1,1]) size = random.choice(cfg.train_sizes) max_size = cfg.train_max_size else: scale, rot, color_scale, do_flip, crop = 1.0, 0.0, np.array([1, 1, 1]), False, np.array([1,1]) size = random.choice(cfg.test_sizes) max_size = cfg.test_max_size crop_bbox_wh = (bbox[2:]*crop).astype(np.uint32) xy_range = img.shape[:2][::-1]-crop_bbox_wh crop_bbox_xywh = np.array([np.random.randint(0,xy_range[0]+1),np.random.randint(0,xy_range[1]+1),crop_bbox_wh[0],crop_bbox_wh[1]]) reshape_size = resize(crop_bbox_xywh[2:], size, max_size) img, trans, inv_trans = generate_patch_image(img, crop_bbox_xywh, 1, rot, do_flip, reshape_size[::-1]) img = np.clip(img * color_scale[None, None, :], 0, 255) return img, trans, inv_trans, rot, do_flip def augmentation_instance_sample(img, bbox, data_split,data,dataname): ori_shape = img.shape[:2][::-1] if getattr(cfg, 'no_aug', False) and data_split == 'train': scale, rot, color_scale, do_flip,size,crop,sample_ratio,sample_prob = 1.0, 0.0, np.array([1, 1, 1]), False, ori_shape, np.array([1,1]), 0,0 size = random.choice(cfg.train_sizes) max_size = cfg.train_max_size elif data_split == 'train': scale, rot, color_scale, do_flip, crop, sample_ratio,sample_prob = get_aug_config(dataname) rot=0 # scale, rot, do_flip, crop = 1.0, 0.0, False, np.array([1,1]) size = random.choice(cfg.train_sizes) max_size = cfg.train_max_size else: scale, rot, color_scale, do_flip, crop,sample_ratio,sample_prob = 1.0, 0.0, np.array([1, 1, 1]), False, np.array([1,1]),0,0 size = random.choice(cfg.test_sizes) max_size = cfg.test_max_size if random.random() < sample_prob: crop_person_number = len(data['bbox']) if random.random() < sample_ratio: if random.random() < 0.6: crop_person_number_sample = 1 else: crop_person_number_sample = np.random.randint(crop_person_number) + 1 else: crop_person_number_sample = crop_person_number sample_ids = np.array( random.sample(list(range(crop_person_number)), crop_person_number_sample)) bbox_xyxy = [] bbox_xyxy = np.stack(data['bbox'],axis=0)[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])]) crop_bbox_xyxy = np.concatenate([leftTop_, rightBottom_]) crop_bbox_xywh = crop_bbox_xyxy.copy() crop_bbox_xywh[2:] = crop_bbox_xywh[2:]-crop_bbox_xywh[:2] crop_bbox_xywh = adjust_bounding_box(crop_bbox_xywh,ori_shape[0],ori_shape[1]) else: crop_bbox_xywh = bbox.copy() reshape_size = resize(crop_bbox_xywh[2:], size, max_size) # try: # reshape_size = resize(crop_bbox_xywh[2:], size, max_size) # except Exception as e: # print(crop_bbox_xywh) # print(size) # print(max_size) # raise e img, trans, inv_trans = generate_patch_image(img, crop_bbox_xywh, 1, rot, do_flip, reshape_size[::-1]) img = np.clip(img * color_scale[None, None, :], 0, 255) return img, trans, inv_trans, rot, do_flip def adjust_bounding_box(input_bbox,image_width, image_height): left_x, left_y, width, height = input_bbox # Calculate original bounding box center original_center_x = left_x + width / 2 original_center_y = left_y + height / 2 # Calculate target aspect ratio target_aspect_ratio = image_width / image_height # Adjust width and height to match target aspect ratio if width / height > target_aspect_ratio: # Bounding box is wider, adjust height new_height = width / target_aspect_ratio new_width = width else: # Bounding box is taller, adjust width new_width = height * target_aspect_ratio new_height = height # Calculate new bounding box center new_center_x = original_center_x new_center_y = original_center_y # Check if the adjusted bounding box is out of the image boundaries if new_center_x - new_width / 2 < 0: # Shift the bounding box to the right to fit within the image new_center_x = new_width / 2 elif new_center_x + new_width / 2 > image_width: # Shift the bounding box to the left to fit within the image new_center_x = image_width - new_width / 2 if new_center_y - new_height / 2 < 0: # Shift the bounding box down to fit within the image new_center_y = new_height / 2 elif new_center_y + new_height / 2 > image_height: # Shift the bounding box up to fit within the image new_center_y = image_height - new_height / 2 # Calculate adjusted left x and left y of the bounding box and convert to integers adjusted_left_x = int(new_center_x - new_width / 2) adjusted_left_y = int(new_center_y - new_height / 2) # Ensure width and height are integers as well adjusted_width = int(new_width) adjusted_height = int(new_height) # Return adjusted bounding box coordinates (left x, left y, width, height) return np.array([adjusted_left_x, adjusted_left_y, adjusted_width, adjusted_height]) def generate_patch_image(cvimg, bbox, scale, rot, do_flip, out_shape): img = cvimg.copy() img_height, img_width, img_channels = img.shape bb_c_x = float(bbox[0] + 0.5 * bbox[2]) bb_c_y = float(bbox[1] + 0.5 * bbox[3]) bb_width = float(bbox[2]) bb_height = float(bbox[3]) if do_flip: img = img[:, ::-1, :] bb_c_x = img_width - bb_c_x - 1 trans = gen_trans_from_patch_cv(bb_c_x, bb_c_y, bb_width, bb_height, out_shape[1], out_shape[0], scale, rot) img_patch = cv2.warpAffine(img, trans, (int(out_shape[1]), int(out_shape[0])), flags=cv2.INTER_LINEAR) img_patch = img_patch.astype(np.float32) inv_trans = gen_trans_from_patch_cv(bb_c_x, bb_c_y, bb_width, bb_height, out_shape[1], out_shape[0], scale, rot, inv=True) return img_patch, trans, inv_trans def rotate_2d(pt_2d, rot_rad): x = pt_2d[0] y = pt_2d[1] sn, cs = np.sin(rot_rad), np.cos(rot_rad) xx = x * cs - y * sn yy = x * sn + y * cs return np.array([xx, yy], dtype=np.float32) def gen_trans_from_patch_cv(c_x, c_y, src_width, src_height, dst_width, dst_height, scale, rot, inv=False): # augment size with scale src_w = src_width * scale src_h = src_height * scale src_center = np.array([c_x, c_y], dtype=np.float32) # augment rotation rot_rad = np.pi * rot / 180 src_downdir = rotate_2d(np.array([0, src_h * 0.5], dtype=np.float32), rot_rad) src_rightdir = rotate_2d(np.array([src_w * 0.5, 0], dtype=np.float32), rot_rad) dst_w = dst_width dst_h = dst_height dst_center = np.array([dst_w * 0.5, dst_h * 0.5], dtype=np.float32) dst_downdir = np.array([0, dst_h * 0.5], dtype=np.float32) dst_rightdir = np.array([dst_w * 0.5, 0], dtype=np.float32) src = np.zeros((3, 2), dtype=np.float32) src[0, :] = src_center src[1, :] = src_center + src_downdir src[2, :] = src_center + src_rightdir dst = np.zeros((3, 2), dtype=np.float32) dst[0, :] = dst_center dst[1, :] = dst_center + dst_downdir dst[2, :] = dst_center + dst_rightdir if inv: trans = cv2.getAffineTransform(np.float32(dst), np.float32(src)) else: trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) trans = trans.astype(np.float32) return trans def process_db_coord_batch_no_valid(joint_img, joint_cam, do_flip, img_shape, flip_pairs, img2bb_trans, rot, src_joints_name, target_joints_name, input_img_shape): joint_img_original = joint_img.copy() joint_img, joint_cam = joint_img.copy(), joint_cam.copy() # flip augmentation if do_flip: joint_cam[:, :, 0] = -joint_cam[:, :, 0] joint_img[:, :, 0] = img_shape[1] - 1 - joint_img[:, :, 0] for pair in flip_pairs: joint_img[:, pair[0], :], joint_img[:, pair[ 1], :] = joint_img[:, pair[1], :].copy( ), joint_img[:, pair[0], :].copy() joint_cam[:, pair[0], :], joint_cam[:, pair[ 1], :] = joint_cam[:, pair[1], :].copy( ), joint_cam[:, pair[0], :].copy() # 3D data rotation augmentation rot_aug_mat = np.array( [[np.cos(np.deg2rad(-rot)), -np.sin(np.deg2rad(-rot)), 0], [np.sin(np.deg2rad(-rot)), np.cos(np.deg2rad(-rot)), 0], [0, 0, 1]], dtype=np.float32) num_p, num_joints, joints_dim = joint_cam.shape joint_cam = joint_cam.reshape(num_p * num_joints, joints_dim) joint_cam[:,:-1] = np.dot(rot_aug_mat, joint_cam[:,:-1].transpose(1, 0)).transpose(1, 0) joint_cam = joint_cam.reshape(num_p, num_joints, joints_dim) # affine transformation joint_img_xy1 = \ np.concatenate((joint_img[:, :, :2], np.ones_like(joint_img[:, :, :1])), 2) joint_img_xy1 = joint_img_xy1.reshape(num_p * num_joints, 3) joint_img[:, :, :2] = np.dot(img2bb_trans, joint_img_xy1.transpose(1, 0)).transpose( 1, 0).reshape(num_p, num_joints, 2) joint_img[:, :, 0] = joint_img[:, :, 0] / input_img_shape[1] * cfg.output_hm_shape[2] joint_img[:, :, 1] = joint_img[:, :, 1] / input_img_shape[0] * cfg.output_hm_shape[1] # check truncation # TODO # remove 3rd joint_trunc = ((joint_img_original[:,:, 0] >= 0) * (joint_img[:,:, 0] >= 0) * (joint_img[:,:, 0] < cfg.output_hm_shape[2]) * \ (joint_img_original[:,:, 1] >= 0) *(joint_img[:,:, 1] >= 0) * (joint_img[:,:, 1] < cfg.output_hm_shape[1]) * \ joint_img[:,:, -1] ).reshape(num_p, -1, 1).astype(np.float32) # transform joints to target db joints joint_img = transform_joint_to_other_db_batch(joint_img, src_joints_name, target_joints_name) joint_cam_wo_ra = transform_joint_to_other_db_batch( joint_cam, src_joints_name, target_joints_name) joint_trunc = transform_joint_to_other_db_batch(joint_trunc, src_joints_name, target_joints_name) # root-alignment, for joint_cam input wo ra joint_cam_ra = joint_cam_wo_ra.copy() joint_cam_ra[:,:,:3] = joint_cam_ra[:,:,:3] - joint_cam_ra[:, smpl_x.root_joint_idx, None, :3] # root-relative joint_cam_ra[:, smpl_x.joint_part[ 'lhand'], :3] = joint_cam_ra[:, smpl_x.joint_part[ 'lhand'], :3] - joint_cam_ra[:, smpl_x.lwrist_idx, None, :3] # left hand root-relative joint_cam_ra[:, smpl_x.joint_part[ 'rhand'], :3] = joint_cam_ra[:, smpl_x.joint_part[ 'rhand'], :3] - joint_cam_ra[:, smpl_x.rwrist_idx, None, :3] # right hand root-relative joint_cam_ra[:, smpl_x. joint_part['face'], :3] = joint_cam_ra[:, smpl_x.joint_part[ 'face'], :3] - joint_cam_ra[:, smpl_x.neck_idx, None, :3] # face root-relative return joint_img, joint_cam_wo_ra, joint_cam_ra, joint_trunc def process_human_model_output_batch_ubody(human_model_param, do_flip, rot, as_smplx, part_valid ): num_person = human_model_param['body_pose'].shape[0] human_model = smpl_x rotation_valid = np.ones((num_person,smpl_x.orig_joint_num), dtype=np.float32) coord_valid = np.ones((num_person,smpl_x.joint_num), dtype=np.float32) # expr_valid = np.ones((num_person), dtype=np.float32) # shape_valid = np.ones((num_person), dtype=np.float32) # root_pose, body_pose, shape, trans = human_model_param['root_pose'], human_model_param['body_pose'], \ # human_model_param['shape'], human_model_param['trans'] if 'smplx_valid' in human_model_param: smplx_valid = human_model_param['smplx_valid'] shape_valid = human_model_param['smplx_valid'] else: smplx_valid = np.ones(num_person, dtype=np.bool8) shape_valid = np.ones(num_person, dtype=np.bool8) if 'expr_valid' in human_model_param: expr_valid = human_model_param['expr_valid'] else: expr_valid = np.ones(num_person, dtype=np.bool8) expr_valid*=smplx_valid if 'face_valid' in human_model_param: face_valid = human_model_param['face_valid'] else: face_valid = np.ones(num_person, dtype=np.bool8) face_valid *= smplx_valid # check lhand valid key exsits if 'lhand_valid' in human_model_param: lhand_valid = human_model_param['lhand_valid'] else: lhand_valid = np.ones(num_person, dtype=np.bool8) lhand_valid*=smplx_valid # check rhand valid key exsits if 'rhand_valid' in human_model_param: rhand_valid = human_model_param['rhand_valid'] else: rhand_valid = np.ones(num_person, dtype=np.bool8) rhand_valid*=smplx_valid # check validation of the smplx parameters if 'body_pose' in human_model_param \ and human_model_param['body_pose'] is not None: root_pose, body_pose = human_model_param['root_pose'], human_model_param['body_pose'] shape, trans = human_model_param['shape'], human_model_param['trans'] root_pose = torch.FloatTensor(root_pose).view(num_person, 1, 3) body_pose = torch.FloatTensor(body_pose).view(num_person, -1, 3) shape = torch.FloatTensor(shape).view(num_person, -1) trans = torch.FloatTensor(trans).view(num_person,-1) else: root_pose = np.zeros((num_person, 3), dtype=np.float32) body_pose = np.zeros((num_person, 3 * len(smpl_x.orig_joint_part['body'])), dtype=np.float32) shape = np.zeros((num_person, 10), dtype=np.float32) trans = np.zeros((num_person, 3), dtype=np.float32) rotation_valid[:, smpl_x.orig_joint_part['body']] = 0 coord_valid[:, smpl_x.joint_part['body']] = 0 body_pose*=smplx_valid[:, None, None] root_pose*=smplx_valid[:, None, None] shape*=smplx_valid[:, None] trans*=smplx_valid[:, None] rotation_valid[:, smpl_x.orig_joint_part['body']]*=smplx_valid[:,None] coord_valid[:, smpl_x.joint_part['body']]*=smplx_valid[:,None] # check validation of the smplx parameters if 'lhand_pose' in human_model_param \ and human_model_param['lhand_pose'] is not None: lhand_pose = human_model_param['lhand_pose'] lhand_pose = torch.FloatTensor(lhand_pose).view(num_person, -1, 3) # lhand_valid = part_valid['lhand'] # rotation_valid[:, smpl_x.orig_joint_part['lhand']]*=lhand_valid[:,None] # coord_valid[:, smpl_x.joint_part['lhand']]*=lhand_valid[:,None] else: lhand_pose = np.zeros((num_person, 3 * len(smpl_x.orig_joint_part['lhand'])), dtype=np.float32) rotation_valid[:, smpl_x.orig_joint_part['lhand']] = 0 coord_valid[:, smpl_x.joint_part['lhand']] = 0 lhand_pose*=lhand_valid[:,None,None] rotation_valid[:, smpl_x.orig_joint_part['lhand']]*=lhand_valid[:,None] coord_valid[:, smpl_x.joint_part['lhand']]*=lhand_valid[:,None] if 'rhand_pose' in human_model_param \ and human_model_param['rhand_pose'] is not None: rhand_pose = human_model_param['rhand_pose'] rhand_pose = torch.FloatTensor(rhand_pose).view(num_person, -1, 3) # rhand_valid = part_valid['rhand'] # rotation_valid[:, smpl_x.orig_joint_part['rhand']]*=rhand_valid[:,None] # coord_valid[:, smpl_x.joint_part['rhand']]*=rhand_valid[:,None] else: rhand_pose = np.zeros((num_person, 3 * len(smpl_x.orig_joint_part['rhand'])), dtype=np.float32) rotation_valid[:, smpl_x.orig_joint_part['rhand']] = 0 coord_valid[:, smpl_x.joint_part['rhand']] = 0 rhand_pose*=rhand_valid[:,None,None] rotation_valid[:, smpl_x.orig_joint_part['rhand']]*=rhand_valid[:,None] coord_valid[:, smpl_x.joint_part['rhand']]*=rhand_valid[:,None] if 'expr' in human_model_param and \ human_model_param['expr'] is not None: expr = human_model_param['expr'] # face_valid = part_valid['face'] # expr_valid = expr_valid*face_valid else: expr = np.zeros((num_person, smpl_x.expr_code_dim), dtype=np.float32) expr_valid = expr_valid*0 expr*=face_valid[:,None] expr = torch.FloatTensor(expr).view(num_person,-1) expr_valid*=face_valid # expr is invalid if face_valid is 0 if 'jaw_pose' in human_model_param and \ human_model_param['jaw_pose'] is not None: jaw_pose = human_model_param['jaw_pose'] # face_valid = part_valid['face'] # rotation_valid[:, smpl_x.orig_joint_part['face']]*=face_valid[:,None] # coord_valid[:, smpl_x.joint_part['face']]*=face_valid[:,None] else: jaw_pose = np.zeros((num_person, 3), dtype=np.float32) rotation_valid[:,smpl_x.orig_joint_part['face']] = 0 coord_valid[:,smpl_x.joint_part['face']] = 0 jaw_pose*=face_valid[:,None] jaw_pose = torch.FloatTensor(jaw_pose).view(num_person, -1, 3) rotation_valid[:, smpl_x.orig_joint_part['face']]*=face_valid[:,None] coord_valid[:, smpl_x.joint_part['face']]*=face_valid[:,None] if 'gender' in human_model_param and \ human_model_param['gender'] is not None: gender = human_model_param['gender'] else: gender = 'neutral' if as_smplx == 'smpl': rotation_valid[:,:] = 0 rotation_valid[:,:21] = 1 expr_valid = expr_valid*0 coord_valid[:,:] = 0 coord_valid[:,smpl_x.joint_part['body']] = 1 root_pose = torch.FloatTensor(root_pose).view(num_person, 1, 3) body_pose = torch.FloatTensor(body_pose).view(num_person, -1, 3) lhand_pose = torch.FloatTensor( lhand_pose).view(num_person, -1, 3) rhand_pose = torch.FloatTensor(rhand_pose).view(num_person, -1, 3) jaw_pose = torch.FloatTensor(jaw_pose).view(num_person, -1, 3) shape = torch.FloatTensor(shape).view(num_person, -1) expr = torch.FloatTensor(expr).view(num_person,-1) trans = torch.FloatTensor(trans).view(num_person,-1) pose = torch.cat((root_pose, body_pose, lhand_pose, rhand_pose, jaw_pose),dim=1) ## so far, data augmentations are not applied yet ## now, apply data augmentations # x,y affine transform, root-relative depth # 3D data rotation augmentation # rot_aug_mat = np.array( # [[np.cos(np.deg2rad(-rot)), -np.sin(np.deg2rad(-rot)), 0], # [np.sin(np.deg2rad(-rot)), # np.cos(np.deg2rad(-rot)), 0], [0, 0, 1]], # dtype=np.float32) # parameters # flip pose parameter (axis-angle) if do_flip: for pair in human_model.orig_flip_pairs: pose[:, pair[0], :], pose[:, pair[1], :] = pose[:, pair[1], :].clone( ), pose[:, pair[0], :].clone() rotation_valid[:,pair[0]], rotation_valid[:,pair[1]] = rotation_valid[:,pair[1]].copy(), rotation_valid[:, pair[0]].copy() pose[:,:, 1:3] *= -1 # multiply -1 to y and z axis of axis-angle # rotate root pose pose = pose.numpy() root_pose = pose[:, human_model.orig_root_joint_idx, :] # for pose_i in range(len(root_pose)): # root_pose_mat = cv2.Rodrigues(root_pose[pose_i])[0] # root_pose[pose_i] = cv2.Rodrigues(np.dot(rot_aug_mat, # root_pose_mat))[0][:, 0] pose[:, human_model.orig_root_joint_idx] = root_pose.reshape(num_person, 3) # change to mean shape if beta is too far from it # shape[(shape.abs() > 3).any(dim=1)] = 0. shape = shape.numpy().reshape(num_person, -1) # shape_valid = shape.sum(-1)!=0 # return results pose = pose.reshape(num_person, -1) expr = expr.numpy().reshape(num_person, -1) return pose, shape, expr, rotation_valid, coord_valid, expr_valid, shape_valid def process_human_model_output_batch_simplify(human_model_param, do_flip, rot, as_smplx, data_name=None ): num_person = human_model_param['body_pose'].shape[0] human_model = smpl_x rotation_valid = np.ones((num_person,smpl_x.orig_joint_num), dtype=np.float32) coord_valid = np.ones((num_person,smpl_x.joint_num), dtype=np.float32) # expr_valid = np.ones((num_person), dtype=np.float32) # shape_valid = np.ones((num_person), dtype=np.float32) # shape, trans = human_model_param['shape'], human_model_param['trans'] # check smplx valid key exsits if 'smplx_valid' in human_model_param: smplx_valid = human_model_param['smplx_valid'] shape_valid = human_model_param['smplx_valid'] else: smplx_valid = np.ones(num_person, dtype=np.bool8) shape_valid = np.ones(num_person, dtype=np.bool8) if 'expr_valid' in human_model_param: expr_valid = human_model_param['expr_valid'] else: expr_valid = np.ones(num_person, dtype=np.bool8) expr_valid*=smplx_valid # check face valid key exsits if 'face_valid' in human_model_param: face_valid = human_model_param['face_valid'] else: face_valid = np.ones(num_person, dtype=np.bool8) face_valid *= smplx_valid # check lhand valid key exsits if 'lhand_valid' in human_model_param: lhand_valid = human_model_param['lhand_valid'] else: lhand_valid = np.ones(num_person, dtype=np.bool8) lhand_valid*=smplx_valid # check rhand valid key exsits if 'rhand_valid' in human_model_param: rhand_valid = human_model_param['rhand_valid'] else: rhand_valid = np.ones(num_person, dtype=np.bool8) rhand_valid*=smplx_valid # check validation of the smplx parameters if 'body_pose' in human_model_param \ and human_model_param['body_pose'] is not None: root_pose, body_pose = human_model_param['root_pose'], human_model_param['body_pose'] shape, trans = human_model_param['shape'], human_model_param['trans'] root_pose = torch.FloatTensor(root_pose).view(num_person, 1, 3) body_pose = torch.FloatTensor(body_pose).view(num_person, -1, 3) shape = torch.FloatTensor(shape).view(num_person, -1) trans = torch.FloatTensor(trans).view(num_person,-1) else: root_pose = np.zeros((num_person, 3), dtype=np.float32) body_pose = np.zeros((num_person, 3 * len(smpl_x.orig_joint_part['body'])), dtype=np.float32) shape = np.zeros((num_person, 10), dtype=np.float32) trans = np.zeros((num_person, 3), dtype=np.float32) rotation_valid[:, smpl_x.orig_joint_part['body']] = 0 coord_valid[:, smpl_x.joint_part['body']] = 0 body_pose*=smplx_valid[:, None, None] root_pose*=smplx_valid[:, None, None] shape*=smplx_valid[:, None] trans*=smplx_valid[:, None] rotation_valid[:, smpl_x.orig_joint_part['body']]*=smplx_valid[:,None] coord_valid[:, smpl_x.joint_part['body']]*=smplx_valid[:,None] if 'lhand_pose' in human_model_param \ and human_model_param['lhand_pose'] is not None: lhand_pose = human_model_param['lhand_pose'] lhand_pose = torch.FloatTensor(lhand_pose).view(num_person, -1, 3) else: lhand_pose = np.zeros((num_person, 3 * len(smpl_x.orig_joint_part['lhand'])), dtype=np.float32) rotation_valid[:, smpl_x.orig_joint_part['lhand']] = 0 coord_valid[:, smpl_x.joint_part['lhand']] = 0 lhand_pose*=lhand_valid[:,None,None] rotation_valid[:, smpl_x.orig_joint_part['lhand']]*=lhand_valid[:,None] coord_valid[:, smpl_x.joint_part['lhand']]*=lhand_valid[:,None] if 'rhand_pose' in human_model_param \ and human_model_param['rhand_pose'] is not None: rhand_pose = human_model_param['rhand_pose'] rhand_pose = torch.FloatTensor(rhand_pose).view(num_person, -1, 3) else: rhand_pose = np.zeros((num_person, 3 * len(smpl_x.orig_joint_part['rhand'])), dtype=np.float32) rotation_valid[:, smpl_x.orig_joint_part['rhand']] = 0 coord_valid[:, smpl_x.joint_part['rhand']] = 0 rhand_pose*=rhand_valid[:,None,None] rotation_valid[:, smpl_x.orig_joint_part['rhand']]*=rhand_valid[:,None] coord_valid[:, smpl_x.joint_part['rhand']]*=rhand_valid[:,None] # face valid > expr valid > face kps valid, but for synbody and bedlam if 'expr' in human_model_param and \ human_model_param['expr'] is not None: expr = human_model_param['expr'] else: expr = np.zeros((num_person, smpl_x.expr_code_dim), dtype=np.float32) expr_valid = expr_valid * 0 expr*=face_valid[:,None] expr = torch.FloatTensor(expr).view(num_person,-1) expr_valid*=face_valid # expr is invalid if face_valid is 0 # for coco and ubody, if face is invalid, jaw pose and face kps2d should be false if 'jaw_pose' in human_model_param and \ human_model_param['jaw_pose'] is not None: jaw_pose = human_model_param['jaw_pose'] else: jaw_pose = np.zeros((num_person, 3), dtype=np.float32) rotation_valid[:,smpl_x.orig_joint_part['face']] = 0 coord_valid[:,smpl_x.joint_part['face']] = 0 # if data_name not in ["BEDLAM"]: # face_valid = face_valid * expr_valid * jaw_pose # else: # # synbody and bedlam exps valid is false but jaw pose and face kps2d is valid # face_valid = face_valid jaw_pose*=face_valid[:,None] jaw_pose = torch.FloatTensor(jaw_pose).view(num_person, -1, 3) rotation_valid[:, smpl_x.orig_joint_part['face']]*=face_valid[:,None] coord_valid[:, smpl_x.joint_part['face']]*=face_valid[:,None] # if data_name not in ["BEDLAM" , "SynBody"]: # coord_valid[:, smpl_x.joint_part['face']] = coord_valid[:, smpl_x.joint_part['face']] * expr_valid[:,None] # expr valid? # coord_valid[:, smpl_x.joint_part['face']] = coord_valid[:, smpl_x.joint_part['face']] * face_valid[:,None] if 'gender' in human_model_param and \ human_model_param['gender'] is not None: gender = human_model_param['gender'] else: gender = 'neutral' if as_smplx == 'smpl': rotation_valid[:,:] = 0 rotation_valid[:,:21] = 1 expr_valid = expr_valid*0 coord_valid[:,:] = 0 coord_valid[:,smpl_x.joint_part['body']] = 1 # print(root_pose.shape, body_pose.shape, lhand_pose.shape, rhand_pose.shape, jaw_pose.shape) pose = torch.cat((root_pose, body_pose, lhand_pose, rhand_pose, jaw_pose),dim=1) ## so far, data augmentations are not applied yet ## now, apply data augmentations # 3D data rotation augmentation # rot_aug_mat = np.array( # [[np.cos(np.deg2rad(-rot)), -np.sin(np.deg2rad(-rot)), 0], # [np.sin(np.deg2rad(-rot)), # np.cos(np.deg2rad(-rot)), 0], [0, 0, 1]], # dtype=np.float32) # parameters # flip pose parameter (axis-angle) if do_flip: for pair in human_model.orig_flip_pairs: pose[:, pair[0], :], pose[:, pair[1], :] = pose[:, pair[1], :].clone( ), pose[:, pair[0], :].clone() rotation_valid[:,pair[0]], rotation_valid[:,pair[1]] = rotation_valid[:,pair[1]].copy(), rotation_valid[:, pair[0]].copy() pose[:,:, 1:3] *= -1 # multiply -1 to y and z axis of axis-angle # rotate root pose pose = pose.numpy() root_pose = pose[:, human_model.orig_root_joint_idx, :] # for pose_i in range(len(root_pose)): # root_pose_mat = cv2.Rodrigues(root_pose[pose_i])[0] # root_pose[pose_i] = cv2.Rodrigues(np.dot(rot_aug_mat, # root_pose_mat))[0][:, 0] pose[:, human_model.orig_root_joint_idx] = root_pose.reshape(num_person, 3) # change to mean shape if beta is too far from it # shape[(shape.abs() > 3).any(dim=1)] = 0. shape = shape.numpy().reshape(num_person, -1) # shape_valid = shape.sum(-1)!=0 # return results pose = pose.reshape(num_person, -1) expr = expr.numpy().reshape(num_person, -1) return pose, shape, expr, rotation_valid, coord_valid, expr_valid, shape_valid def load_obj(file_name): v = [] obj_file = open(file_name) for line in obj_file: words = line.split(' ') if words[0] == 'v': x, y, z = float(words[1]), float(words[2]), float(words[3]) v.append(np.array([x, y, z])) return np.stack(v) def load_ply(file_name): plydata = PlyData.read(file_name) x = plydata['vertex']['x'] y = plydata['vertex']['y'] z = plydata['vertex']['z'] v = np.stack((x, y, z), 1) return v def resize_bbox(bbox, scale=1.2): if isinstance(bbox, list): x1, y1, x2, y2 = bbox[0], bbox[1], bbox[2], bbox[3] else: x1, y1, x2, y2 = bbox x_center = (x1 + x2) / 2.0 y_center = (y1 + y2) / 2.0 x_size, y_size = x2 - x1, y2 - y1 x1_resize = x_center - x_size / 2.0 * scale x2_resize = x_center + x_size / 2.0 * scale y1_resize = y_center - y_size / 2.0 * scale y2_resize = y_center + y_size / 2.0 * scale bbox[0], bbox[1], bbox[2], bbox[ 3] = x1_resize, y1_resize, x2_resize, y2_resize return bbox