import os import os.path as osp from glob import glob import numpy as np from config import cfg import copy import json import pickle import cv2 import torch from pycocotools.coco import COCO from utils.human_models import smpl_x from utils.preprocessing import load_img, sanitize_bbox, process_bbox, augmentation, process_db_coord, \ process_human_model_output, load_ply, load_obj from utils.transforms import rigid_align import tqdm import random from humandata import Cache class AGORA(torch.utils.data.Dataset): def __init__(self, transform, data_split): self.transform = transform self.data_split = data_split if getattr(cfg, 'eval_on_train', False): self.data_split = 'eval_train' print("Evaluate on train set.") self.data_path = osp.join(cfg.data_dir, 'AGORA', 'data') self.save_idx = 0 self.resolution = (2160, 3840) # height, width. one of (720, 1280) and (2160, 3840) if cfg.agora_benchmark == 'agora_model_test' or cfg.agora_benchmark == 'test_only': self.test_set = 'test' else: self.test_set = 'val' # val, test # AGORA joint set self.joint_set = { 'joint_num': 127, 'joints_name': \ ('Pelvis', 'L_Hip', 'R_Hip', 'Spine_1', 'L_Knee', 'R_Knee', 'Spine_2', 'L_Ankle', 'R_Ankle', 'Spine_3', 'L_Foot', 'R_Foot', 'Neck', 'L_Collar', 'R_Collar', 'Head', 'L_Shoulder', 'R_Shoulder', 'L_Elbow', 'R_Elbow', 'L_Wrist', 'R_Wrist', # body 'Jaw', 'L_Eye_SMPLH', 'R_Eye_SMPLH', # SMPLH 'L_Index_1', 'L_Index_2', 'L_Index_3', 'L_Middle_1', 'L_Middle_2', 'L_Middle_3', 'L_Pinky_1', 'L_Pinky_2', 'L_Pinky_3', 'L_Ring_1', 'L_Ring_2', 'L_Ring_3', 'L_Thumb_1', 'L_Thumb_2', 'L_Thumb_3', # fingers 'R_Index_1', 'R_Index_2', 'R_Index_3', 'R_Middle_1', 'R_Middle_2', 'R_Middle_3', 'R_Pinky_1', 'R_Pinky_2', 'R_Pinky_3', 'R_Ring_1', 'R_Ring_2', 'R_Ring_3', 'R_Thumb_1', 'R_Thumb_2', 'R_Thumb_3', # fingers 'Nose', 'R_Eye', 'L_Eye', 'R_Ear', 'L_Ear', # face in body 'L_Big_toe', 'L_Small_toe', 'L_Heel', 'R_Big_toe', 'R_Small_toe', 'R_Heel', # feet 'L_Thumb_4', 'L_Index_4', 'L_Middle_4', 'L_Ring_4', 'L_Pinky_4', # finger tips 'R_Thumb_4', 'R_Index_4', 'R_Middle_4', 'R_Ring_4', 'R_Pinky_4', # finger tips *['Face_' + str(i) for i in range(5, 56)] # face ), 'flip_pairs': \ ((1, 2), (4, 5), (7, 8), (10, 11), (13, 14), (16, 17), (18, 19), (20, 21), # body (23, 24), # SMPLH (25, 40), (26, 41), (27, 42), (28, 43), (29, 44), (30, 45), (31, 46), (32, 47), (33, 48), (34, 49), (35, 50), (36, 51), (37, 52), (38, 53), (39, 54), # fingers (56, 57), (58, 59), # face in body (60, 63), (61, 64), (62, 65), # feet (66, 71), (67, 72), (68, 73), (69, 74), (70, 75), # fingertips (76, 85), (77, 84), (78, 83), (79, 82), (80, 81), # face eyebrow (90, 94), (91, 93), # face below nose (95, 104), (96, 103), (97, 102), (98, 101), (99, 106), (100, 105), # face eyes (107, 113), (108, 112), (109, 111), (114, 118), (115, 117), # face mouth (119, 123), (120, 122), (124, 126) # face lip ) } self.joint_set['joint_part'] = { 'body': list(range(self.joint_set['joints_name'].index('Pelvis'), self.joint_set['joints_name'].index('R_Eye_SMPLH') + 1)) + list( range(self.joint_set['joints_name'].index('Nose'), self.joint_set['joints_name'].index('R_Heel') + 1)), 'lhand': list(range(self.joint_set['joints_name'].index('L_Index_1'), self.joint_set['joints_name'].index('L_Thumb_3') + 1)) + list( range(self.joint_set['joints_name'].index('L_Thumb_4'), self.joint_set['joints_name'].index('L_Pinky_4') + 1)), 'rhand': list(range(self.joint_set['joints_name'].index('R_Index_1'), self.joint_set['joints_name'].index('R_Thumb_3') + 1)) + list( range(self.joint_set['joints_name'].index('R_Thumb_4'), self.joint_set['joints_name'].index('R_Pinky_4') + 1)), 'face': list(range(self.joint_set['joints_name'].index('Face_5'), self.joint_set['joints_name'].index('Face_55') + 1))} self.joint_set['root_joint_idx'] = self.joint_set['joints_name'].index('Pelvis') self.joint_set['lwrist_idx'] = self.joint_set['joints_name'].index('L_Wrist') self.joint_set['rwrist_idx'] = self.joint_set['joints_name'].index('R_Wrist') self.joint_set['neck_idx'] = self.joint_set['joints_name'].index('Neck') # self.datalist = self.load_data() # load data or cache self.use_cache = getattr(cfg, 'use_cache', False) if 'train'in self.data_split or (self.data_split == 'test' and self.test_set == 'val'): if 'train' in self.data_split: if getattr(cfg, 'agora_fix_betas', False): assert getattr(cfg, 'agora_fix_global_orient_transl') self.annot_path_cache = osp.join(cfg.data_dir, 'cache', f'AGORA_{self.data_split}_fix_betas.npz') elif getattr(cfg, 'agora_fix_global_orient_transl', False): self.annot_path_cache = osp.join(cfg.data_dir, 'cache', f'AGORA_{self.data_split}_fix_global_orient_transl.npz') else: self.annot_path_cache = osp.join(cfg.data_dir, 'cache', f'AGORA_{self.data_split}.npz') else: if getattr(cfg, 'agora_fix_betas', False): assert getattr(cfg, 'agora_fix_global_orient_transl') self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'AGORA_validation_fix_betas.npz') elif getattr(cfg, 'agora_fix_global_orient_transl', False): self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'AGORA_validation_fix_global_orient_transl.npz') else: self.annot_path_cache = osp.join(cfg.data_dir, 'cache', 'AGORA_validation.npz') if self.use_cache and osp.isfile(self.annot_path_cache): print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}') datalist = Cache(self.annot_path_cache) assert datalist.data_strategy == getattr(cfg, 'data_strategy', None), \ f'Cache data strategy {datalist.data_strategy} does not match current data strategy ' \ f'{getattr(cfg, "data_strategy", None)}' self.datalist = datalist else: if self.use_cache: print(f'[{self.__class__.__name__}] Cache not found, generating cache...') self.datalist = self.load_data() if self.use_cache: print(f'[{self.__class__.__name__}] Caching datalist to {self.annot_path_cache}...') Cache.save( self.annot_path_cache, self.datalist, data_strategy=getattr(cfg, 'data_strategy', None) ) else: # test self.datalist = self.load_data() def load_data(self): datalist = [] if 'train' in self.data_split or (self.data_split == 'test' and self.test_set == 'val'): print('dataset settings:') print('agora_fix_betas', getattr(cfg, 'agora_fix_betas', False)) print('agora_fix_global_orient_transl', getattr(cfg, 'agora_fix_global_orient_transl', False)) print('agora_valid_root_pose', getattr(cfg, 'agora_valid_root_pose', False)) if 'train' in self.data_split: if getattr(cfg, 'agora_fix_betas', False): assert getattr(cfg, 'agora_fix_global_orient_transl') db = COCO(osp.join(self.data_path, 'AGORA_train_fix_betas.json')) elif getattr(cfg, 'agora_fix_global_orient_transl', False): db = COCO(osp.join(self.data_path, 'AGORA_train_fix_global_orient_transl.json')) else: db = COCO(osp.join(self.data_path, 'AGORA_train.json')) else: if getattr(cfg, 'agora_fix_betas', False): assert getattr(cfg, 'agora_fix_global_orient_transl') db = COCO(osp.join(self.data_path, 'AGORA_validation_fix_betas.json')) elif getattr(cfg, 'agora_fix_global_orient_transl', False): db = COCO(osp.join(self.data_path, 'AGORA_validation_fix_global_orient_transl.json')) else: db = COCO(osp.join(self.data_path, 'AGORA_validation.json')) i = 0 for aid in tqdm.tqdm(list(db.anns.keys())): i += 1 if self.data_split == 'train' and i % getattr(cfg, 'AGORA_train_sample_interval', 1) != 0: continue ann = db.anns[aid] image_id = ann['image_id'] img = db.loadImgs(image_id)[0] if not ann['is_valid']: continue joints_2d_path = osp.join(self.data_path, ann['smplx_joints_2d_path']) joints_3d_path = osp.join(self.data_path, ann['smplx_joints_3d_path']) verts_path = osp.join(self.data_path, ann['smplx_verts_path']) smplx_param_path = osp.join(self.data_path, ann['smplx_param_path']) kid = ann['kid'] gender = ann['gender'] if not osp.exists(smplx_param_path): print(smplx_param_path) if self.resolution == (720, 1280): img_shape = self.resolution img_path = osp.join(self.data_path, img['file_name_1280x720']) # convert to current resolution bbox = np.array(ann['bbox']).reshape(2, 2) bbox[:, 0] = bbox[:, 0] / 3840 * 1280 bbox[:, 1] = bbox[:, 1] / 2160 * 720 bbox = bbox.reshape(4) if hasattr(cfg, 'bbox_ratio'): bbox_ratio = cfg.bbox_ratio * 0.833 # agora preprocess is giving 1.2 box padding else: bbox_ratio = 1.25 bbox = process_bbox(bbox, img_shape[1], img_shape[0], ratio=bbox_ratio) if bbox is None: continue lhand_bbox = np.array(ann['lhand_bbox']).reshape(2, 2) lhand_bbox[:, 0] = lhand_bbox[:, 0] / 3840 * 1280 lhand_bbox[:, 1] = lhand_bbox[:, 1] / 2160 * 720 lhand_bbox = lhand_bbox.reshape(4) lhand_bbox = sanitize_bbox(lhand_bbox, img_shape[1], img_shape[0]) if lhand_bbox is not None: lhand_bbox[2:] += lhand_bbox[:2] # xywh -> xyxy rhand_bbox = np.array(ann['rhand_bbox']).reshape(2, 2) rhand_bbox[:, 0] = rhand_bbox[:, 0] / 3840 * 1280 rhand_bbox[:, 1] = rhand_bbox[:, 1] / 2160 * 720 rhand_bbox = rhand_bbox.reshape(4) rhand_bbox = sanitize_bbox(rhand_bbox, img_shape[1], img_shape[0]) if rhand_bbox is not None: rhand_bbox[2:] += rhand_bbox[:2] # xywh -> xyxy face_bbox = np.array(ann['face_bbox']).reshape(2, 2) face_bbox[:, 0] = face_bbox[:, 0] / 3840 * 1280 face_bbox[:, 1] = face_bbox[:, 1] / 2160 * 720 face_bbox = face_bbox.reshape(4) face_bbox = sanitize_bbox(face_bbox, img_shape[1], img_shape[0]) if face_bbox is not None: face_bbox[2:] += face_bbox[:2] # xywh -> xyxy data_dict = {'img_path': img_path, 'img_shape': img_shape, 'bbox': bbox, 'lhand_bbox': lhand_bbox, 'rhand_bbox': rhand_bbox, 'face_bbox': face_bbox, 'joints_2d_path': joints_2d_path, 'joints_3d_path': joints_3d_path, 'verts_path': verts_path, 'smplx_param_path': smplx_param_path, 'ann_id': str(aid), 'kid': kid, 'gender': gender} datalist.append(data_dict) elif self.resolution == (2160, 3840): # use cropped and resized images. loading 4K images in pytorch dataloader takes too much time... img_path = osp.join(self.data_path, '3840x2160', img['file_name_3840x2160'].split('/')[-2] + '_crop', img['file_name_3840x2160'].split('/')[-1][:-4] + '_ann_id_' + str(aid) + '.png') json_path = osp.join(self.data_path, '3840x2160', img['file_name_3840x2160'].split('/')[-2] + '_crop', img['file_name_3840x2160'].split('/')[-1][:-4] + '_ann_id_' + str( aid) + '.json') if not osp.isfile(json_path): continue with open(json_path) as f: crop_resize_info = json.load(f) img2bb_trans_from_orig = np.array(crop_resize_info['img2bb_trans'], dtype=np.float32) resized_height, resized_width = crop_resize_info['resized_height'], crop_resize_info[ 'resized_width'] img_shape = (resized_height, resized_width) bbox = np.array([0, 0, resized_width, resized_height], dtype=np.float32) # transform from original image to crop_and_resize image lhand_bbox = np.array(ann['lhand_bbox']).reshape(2, 2) lhand_bbox[1] += lhand_bbox[0] # xywh -> xyxy lhand_bbox = np.dot(img2bb_trans_from_orig, np.concatenate((lhand_bbox, np.ones_like(lhand_bbox[:, :1])), 1).transpose(1, 0)).transpose( 1, 0) lhand_bbox[1] -= lhand_bbox[0] # xyxy -> xywh lhand_bbox = lhand_bbox.reshape(4) lhand_bbox = sanitize_bbox(lhand_bbox, self.resolution[1], self.resolution[0]) if lhand_bbox is not None: lhand_bbox[2:] += lhand_bbox[:2] # xywh -> xyxy # transform from original image to crop_and_resize image rhand_bbox = np.array(ann['rhand_bbox']).reshape(2, 2) rhand_bbox[1] += rhand_bbox[0] # xywh -> xyxy rhand_bbox = np.dot(img2bb_trans_from_orig, np.concatenate((rhand_bbox, np.ones_like(rhand_bbox[:, :1])), 1).transpose(1, 0)).transpose( 1, 0) rhand_bbox[1] -= rhand_bbox[0] # xyxy -> xywh rhand_bbox = rhand_bbox.reshape(4) rhand_bbox = sanitize_bbox(rhand_bbox, self.resolution[1], self.resolution[0]) if rhand_bbox is not None: rhand_bbox[2:] += rhand_bbox[:2] # xywh -> xyxy # transform from original image to crop_and_resize image face_bbox = np.array(ann['face_bbox']).reshape(2, 2) face_bbox[1] += face_bbox[0] # xywh -> xyxy face_bbox = np.dot(img2bb_trans_from_orig, np.concatenate((face_bbox, np.ones_like(face_bbox[:, :1])), 1).transpose(1, 0)).transpose( 1, 0) face_bbox[1] -= face_bbox[0] # xyxy -> xywh face_bbox = face_bbox.reshape(4) face_bbox = sanitize_bbox(face_bbox, self.resolution[1], self.resolution[0]) if face_bbox is not None: face_bbox[2:] += face_bbox[:2] # xywh -> xyxy data_dict = {'img_path': img_path, 'img_shape': img_shape, 'bbox': bbox, 'lhand_bbox': lhand_bbox, 'rhand_bbox': rhand_bbox, 'face_bbox': face_bbox, 'img2bb_trans_from_orig': img2bb_trans_from_orig, 'joints_2d_path': joints_2d_path, 'joints_3d_path': joints_3d_path, 'verts_path': verts_path, 'smplx_param_path': smplx_param_path, 'ann_id': str(aid), 'kid': kid, 'gender': gender} datalist.append(data_dict) print('[AGORA train] original size:', len(db.anns.keys()), '. Sample interval:', getattr(cfg, 'AGORA_train_sample_interval', 1), '. Sampled size:', len(datalist)) elif self.data_split == 'test' and self.test_set == 'test': with open(osp.join(self.data_path, 'AGORA_test_bbox.json')) as f: bboxs = json.load(f) for filename in tqdm.tqdm(bboxs.keys()): if self.resolution == (720, 1280): img_path = osp.join(self.data_path, 'test', filename) img_shape = self.resolution person_num = len(bboxs[filename]) for pid in range(person_num): # change bbox from (2160,3840) to target resoution bbox = np.array(bboxs[filename][pid]['bbox']).reshape(2, 2) bbox[:, 0] = bbox[:, 0] / 3840 * 1280 bbox[:, 1] = bbox[:, 1] / 2160 * 720 bbox = bbox.reshape(4) if hasattr(cfg, 'bbox_ratio'): bbox_ratio = cfg.bbox_ratio * 0.833 # agora preprocess is giving 1.2 box padding else: bbox_ratio = 1.25 bbox = process_bbox(bbox, img_shape[1], img_shape[0], ratio=bbox_ratio) if bbox is None: continue datalist.append({'img_path': img_path, 'img_shape': img_shape, 'bbox': bbox, 'person_idx': pid}) elif self.resolution == (2160, 3840): # use cropped and resized images. loading 4K images in pytorch dataloader takes too much time... person_num = len(bboxs[filename]) for pid in range(person_num): img_path = osp.join(self.data_path, '3840x2160', 'test_crop', filename[:-4] + '_pid_' + str(pid) + '.png') json_path = osp.join(self.data_path, '3840x2160', 'test_crop', filename[:-4] + '_pid_' + str(pid) + '.json') if not osp.isfile(json_path): continue with open(json_path) as f: crop_resize_info = json.load(f) img2bb_trans_from_orig = np.array(crop_resize_info['img2bb_trans'], dtype=np.float32) resized_height, resized_width = crop_resize_info['resized_height'], crop_resize_info[ 'resized_width'] img_shape = (resized_height, resized_width) bbox = np.array([0, 0, resized_width, resized_height], dtype=np.float32) datalist.append({'img_path': img_path, 'img_shape': img_shape, 'img2bb_trans_from_orig': img2bb_trans_from_orig, 'bbox': bbox, 'person_idx': pid}) if (getattr(cfg, 'data_strategy', None) == 'balance' and self.data_split == 'train') or \ self.data_split == 'eval_train': print(f"[Agora] Using [balance] strategy with datalist shuffled...") random.seed(2023) random.shuffle(datalist) if self.data_split == "eval_train": return datalist[:10000] return datalist def process_hand_face_bbox(self, bbox, do_flip, img_shape, img2bb_trans): if bbox is None: bbox = np.array([0, 0, 1, 1], dtype=np.float32).reshape(2, 2) # dummy value bbox_valid = float(False) # dummy value else: # reshape to top-left (x,y) and bottom-right (x,y) bbox = bbox.reshape(2, 2) # flip augmentation if do_flip: bbox[:, 0] = img_shape[1] - bbox[:, 0] - 1 bbox[0, 0], bbox[1, 0] = bbox[1, 0].copy(), bbox[0, 0].copy() # xmin <-> xmax swap # make four points of the bbox bbox = bbox.reshape(4).tolist() xmin, ymin, xmax, ymax = bbox bbox = np.array([[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]], dtype=np.float32).reshape(4, 2) # affine transformation (crop, rotation, scale) bbox_xy1 = np.concatenate((bbox, np.ones_like(bbox[:, :1])), 1) bbox = np.dot(img2bb_trans, bbox_xy1.transpose(1, 0)).transpose(1, 0)[:, :2] bbox[:, 0] = bbox[:, 0] / cfg.input_img_shape[1] * cfg.output_hm_shape[2] bbox[:, 1] = bbox[:, 1] / cfg.input_img_shape[0] * cfg.output_hm_shape[1] # make box a rectangle without rotation xmin = np.min(bbox[:, 0]); xmax = np.max(bbox[:, 0]); ymin = np.min(bbox[:, 1]); ymax = np.max(bbox[:, 1]); bbox = np.array([xmin, ymin, xmax, ymax], dtype=np.float32) bbox_valid = float(True) bbox = bbox.reshape(2, 2) return bbox, bbox_valid def __len__(self): return len(self.datalist) def __getitem__(self, idx): data = copy.deepcopy(self.datalist[idx]) img_path, img_shape, bbox = data['img_path'], data['img_shape'], data['bbox'] # image load img = load_img(img_path) # affine transform img, img2bb_trans, bb2img_trans, rot, do_flip = augmentation(img, bbox, self.data_split) img = self.transform(img.astype(np.float32)) / 255. if self.data_split == 'train': # gt load with open(data['joints_2d_path']) as f: joint_img = np.array(json.load(f)).reshape(-1, 2) if self.resolution == (2160, 3840): joint_img[:, :2] = np.dot(data['img2bb_trans_from_orig'], np.concatenate((joint_img, np.ones_like(joint_img[:, :1])), 1).transpose( 1, 0)).transpose(1, 0) # transform from original image to crop_and_resize image joint_img[:, 0] = joint_img[:, 0] / 3840 * self.resolution[1] joint_img[:, 1] = joint_img[:, 1] / 2160 * self.resolution[0] with open(data['joints_3d_path']) as f: joint_cam = np.array(json.load(f)).reshape(-1, 3) with open(data['smplx_param_path'], 'rb') as f: smplx_param = pickle.load(f, encoding='latin1') # hand and face bbox transform lhand_bbox, rhand_bbox, face_bbox = data['lhand_bbox'], data['rhand_bbox'], data['face_bbox'] lhand_bbox, lhand_bbox_valid = self.process_hand_face_bbox(lhand_bbox, do_flip, img_shape, img2bb_trans) rhand_bbox, rhand_bbox_valid = self.process_hand_face_bbox(rhand_bbox, do_flip, img_shape, img2bb_trans) face_bbox, face_bbox_valid = self.process_hand_face_bbox(face_bbox, do_flip, img_shape, img2bb_trans) if do_flip: lhand_bbox, rhand_bbox = rhand_bbox, lhand_bbox lhand_bbox_valid, rhand_bbox_valid = rhand_bbox_valid, lhand_bbox_valid lhand_bbox_center = (lhand_bbox[0] + lhand_bbox[1]) / 2.; rhand_bbox_center = (rhand_bbox[0] + rhand_bbox[1]) / 2.; face_bbox_center = (face_bbox[0] + face_bbox[1]) / 2. lhand_bbox_size = lhand_bbox[1] - lhand_bbox[0]; rhand_bbox_size = rhand_bbox[1] - rhand_bbox[0]; face_bbox_size = face_bbox[1] - face_bbox[0]; # coordinates joint_cam = joint_cam - joint_cam[self.joint_set['root_joint_idx'], None, :] # root-relative joint_cam[self.joint_set['joint_part']['lhand'], :] = joint_cam[self.joint_set['joint_part']['lhand'], :] - joint_cam[self.joint_set['lwrist_idx'], None, :] # left hand root-relative joint_cam[self.joint_set['joint_part']['rhand'], :] = joint_cam[self.joint_set['joint_part']['rhand'], :] - joint_cam[self.joint_set['rwrist_idx'], None, :] # right hand root-relative joint_cam[self.joint_set['joint_part']['face'], :] = joint_cam[self.joint_set['joint_part']['face'], :] - joint_cam[self.joint_set['neck_idx'], None, :] # face root-relative joint_img = np.concatenate((joint_img[:, :2], joint_cam[:, 2:]), 1) # x, y, depth joint_img[self.joint_set['joint_part']['body'], 2] = (joint_cam[self.joint_set['joint_part'][ 'body'], 2].copy() / ( cfg.body_3d_size / 2) + 1) / 2. * \ cfg.output_hm_shape[0] # body depth discretize joint_img[self.joint_set['joint_part']['lhand'], 2] = (joint_cam[self.joint_set['joint_part'][ 'lhand'], 2].copy() / ( cfg.hand_3d_size / 2) + 1) / 2. * \ cfg.output_hm_shape[0] # left hand depth discretize joint_img[self.joint_set['joint_part']['rhand'], 2] = (joint_cam[self.joint_set['joint_part'][ 'rhand'], 2].copy() / ( cfg.hand_3d_size / 2) + 1) / 2. * \ cfg.output_hm_shape[0] # right hand depth discretize joint_img[self.joint_set['joint_part']['face'], 2] = (joint_cam[self.joint_set['joint_part'][ 'face'], 2].copy() / ( cfg.face_3d_size / 2) + 1) / 2. * \ cfg.output_hm_shape[0] # face depth discretize joint_valid = np.ones_like(joint_img[:, :1]) # alr ra when passed into this function joint_img, joint_cam_ra, _, joint_valid, joint_trunc = process_db_coord(joint_img, joint_cam, joint_valid, do_flip, img_shape, self.joint_set['flip_pairs'], img2bb_trans, rot, self.joint_set['joints_name'], smpl_x.joints_name) # reverse ra joint_cam_wo_ra = joint_cam_ra.copy() joint_cam_wo_ra[smpl_x.joint_part['lhand'], :] = joint_cam_wo_ra[smpl_x.joint_part['lhand'], :] \ + joint_cam_wo_ra[smpl_x.lwrist_idx, None, :] # left hand root-relative joint_cam_wo_ra[smpl_x.joint_part['rhand'], :] = joint_cam_wo_ra[smpl_x.joint_part['rhand'], :] \ + joint_cam_wo_ra[smpl_x.rwrist_idx, None, :] # right hand root-relative joint_cam_wo_ra[smpl_x.joint_part['face'], :] = joint_cam_wo_ra[smpl_x.joint_part['face'], :] \ + joint_cam_wo_ra[smpl_x.neck_idx, None,: ] # face root-relative # smplx parameters root_pose = np.array(smplx_param['global_orient'], dtype=np.float32).reshape( -1) # rotation to world coordinate body_pose = np.array(smplx_param['body_pose'], dtype=np.float32).reshape(-1) # use adapted shape for adults if getattr(cfg, 'agora_fix_betas', False) and not data['kid']: shape = np.array(smplx_param['betas_neutral'], dtype=np.float32).reshape(-1)[:10] else: shape = np.array(smplx_param['betas'], dtype=np.float32).reshape(-1)[:10] # bug? lhand_pose = np.array(smplx_param['left_hand_pose'], dtype=np.float32).reshape(-1) rhand_pose = np.array(smplx_param['right_hand_pose'], dtype=np.float32).reshape(-1) jaw_pose = np.array(smplx_param['jaw_pose'], dtype=np.float32).reshape(-1) expr = np.array(smplx_param['expression'], dtype=np.float32).reshape(-1) trans = np.array(smplx_param['transl'], dtype=np.float32).reshape(-1) # translation to world coordinate cam_param = {'focal': cfg.focal, 'princpt': cfg.princpt} # put random camera paraemter as we do not use coordinates from smplx parameters smplx_param = {'root_pose': root_pose, 'body_pose': body_pose, 'shape': shape, 'lhand_pose': lhand_pose, 'lhand_valid': True, 'rhand_pose': rhand_pose, 'rhand_valid': True, 'jaw_pose': jaw_pose, 'expr': expr, 'face_valid': True, 'trans': trans} _, _, _, smplx_pose, smplx_shape, smplx_expr, smplx_pose_valid, _, smplx_expr_valid, _ = process_human_model_output( smplx_param, cam_param, do_flip, img_shape, img2bb_trans, rot, 'smplx') smplx_pose_valid = np.tile(smplx_pose_valid[:, None], (1, 3)).reshape(-1) if not getattr(cfg, 'agora_valid_root_pose', False): smplx_pose_valid[:3] = 0 # global orient of the provided parameter is a rotation to world coordinate system. I want camera coordinate system. smplx_shape_valid = True inputs = {'img': img} targets = {'joint_img': joint_img, 'joint_cam': joint_cam_wo_ra, #from annot 'smplx_joint_img': joint_img, 'smplx_joint_cam': joint_cam_ra, #_smplx_joint_cam, # from smplx param w/ ra 'smplx_pose': smplx_pose, 'smplx_shape': smplx_shape, 'smplx_expr': smplx_expr, 'lhand_bbox_center': lhand_bbox_center, 'lhand_bbox_size': lhand_bbox_size, 'rhand_bbox_center': rhand_bbox_center, 'rhand_bbox_size': rhand_bbox_size, 'face_bbox_center': face_bbox_center, 'face_bbox_size': face_bbox_size} meta_info = {'joint_valid': joint_valid, 'joint_trunc': joint_trunc, 'smplx_joint_valid': joint_valid, 'smplx_joint_trunc': joint_trunc, 'smplx_pose_valid': smplx_pose_valid, 'smplx_shape_valid': float(smplx_shape_valid), 'smplx_expr_valid': float(smplx_expr_valid), 'is_3D': float(True), 'lhand_bbox_valid': lhand_bbox_valid, 'rhand_bbox_valid': rhand_bbox_valid, 'face_bbox_valid': face_bbox_valid} return inputs, targets, meta_info else: # load crop and resize information (for the 4K setting) if self.resolution == (2160, 3840): img2bb_trans = np.dot( np.concatenate((img2bb_trans, np.array([0, 0, 1], dtype=np.float32).reshape(1, 3))), np.concatenate((data['img2bb_trans_from_orig'], np.array([0, 0, 1], dtype=np.float32).reshape(1, 3))) ) bb2img_trans = np.linalg.inv(img2bb_trans)[:2, :] img2bb_trans = img2bb_trans[:2, :] if self.test_set == 'val': # gt load with open(data['verts_path']) as f: verts = np.array(json.load(f)).reshape(-1, 3) with open(data['smplx_param_path'], 'rb') as f: smplx_param = pickle.load(f, encoding='latin1') transl = np.array(smplx_param['transl'], dtype=np.float32).reshape(-1) inputs = {'img': img} targets = {'smplx_mesh_cam': verts} meta_info = {'bb2img_trans': bb2img_trans, 'img_path': img_path, 'gt_smplx_transl':transl} else: inputs = {'img': img} targets = {'smplx_mesh_cam': np.zeros((smpl_x.vertex_num, 3), dtype=np.float32)} # dummy vertex meta_info = {'bb2img_trans': bb2img_trans, 'img_path': img_path} return inputs, targets, meta_info def evaluate(self, outs, cur_sample_idx): annots = self.datalist sample_num = len(outs) eval_result = {'pa_mpvpe_all': [], 'pa_mpvpe_l_hand': [], 'pa_mpvpe_r_hand': [], 'pa_mpvpe_hand': [], 'pa_mpvpe_face': [], 'mpvpe_all': [], 'mpvpe_l_hand': [], 'mpvpe_r_hand': [], 'mpvpe_hand': [], 'mpvpe_face': []} vis = getattr(cfg, 'vis', False) vis_save_dir = cfg.vis_dir if getattr(cfg, 'vis', False): import csv csv_file = f'{cfg.vis_dir}/agora_smplx_error.csv' file = open(csv_file, 'a', newline='') writer = csv.writer(file) for n in range(sample_num): annot = annots[cur_sample_idx + n] out = outs[n] mesh_gt = out['smplx_mesh_cam_target'] mesh_out = out['smplx_mesh_cam'] # MPVPE from all vertices mesh_out_align = mesh_out - np.dot(smpl_x.J_regressor, mesh_out)[smpl_x.J_regressor_idx['pelvis'], None, :] + np.dot(smpl_x.J_regressor, mesh_gt)[smpl_x.J_regressor_idx['pelvis'], None, :] mpvpe_all = np.sqrt(np.sum((mesh_out_align - mesh_gt) ** 2, 1)).mean() * 1000 eval_result['mpvpe_all'].append(mpvpe_all) mesh_out_align = rigid_align(mesh_out, mesh_gt) pa_mpvpe_all = np.sqrt(np.sum((mesh_out_align - mesh_gt) ** 2, 1)).mean() * 1000 eval_result['pa_mpvpe_all'].append(pa_mpvpe_all) # MPVPE from hand vertices mesh_gt_lhand = mesh_gt[smpl_x.hand_vertex_idx['left_hand'], :] mesh_out_lhand = mesh_out[smpl_x.hand_vertex_idx['left_hand'], :] mesh_gt_rhand = mesh_gt[smpl_x.hand_vertex_idx['right_hand'], :] mesh_out_rhand = mesh_out[smpl_x.hand_vertex_idx['right_hand'], :] mesh_out_lhand_align = mesh_out_lhand - np.dot(smpl_x.J_regressor, mesh_out)[ smpl_x.J_regressor_idx['lwrist'], None, :] + np.dot( smpl_x.J_regressor, mesh_gt)[smpl_x.J_regressor_idx['lwrist'], None, :] mesh_out_rhand_align = mesh_out_rhand - np.dot(smpl_x.J_regressor, mesh_out)[ smpl_x.J_regressor_idx['rwrist'], None, :] + np.dot( smpl_x.J_regressor, mesh_gt)[smpl_x.J_regressor_idx['rwrist'], None, :] eval_result['mpvpe_l_hand'].append(np.sqrt( np.sum((mesh_out_lhand_align - mesh_gt_lhand) ** 2, 1)).mean() * 1000) eval_result['mpvpe_r_hand'].append(np.sqrt( np.sum((mesh_out_rhand_align - mesh_gt_rhand) ** 2, 1)).mean() * 1000) eval_result['mpvpe_hand'].append((np.sqrt( np.sum((mesh_out_lhand_align - mesh_gt_lhand) ** 2, 1)).mean() * 1000 + np.sqrt( np.sum((mesh_out_rhand_align - mesh_gt_rhand) ** 2, 1)).mean() * 1000) / 2.) mesh_out_lhand_align = rigid_align(mesh_out_lhand, mesh_gt_lhand) mesh_out_rhand_align = rigid_align(mesh_out_rhand, mesh_gt_rhand) eval_result['pa_mpvpe_l_hand'].append(np.sqrt( np.sum((mesh_out_lhand_align - mesh_gt_lhand) ** 2, 1)).mean() * 1000) eval_result['pa_mpvpe_r_hand'].append(np.sqrt( np.sum((mesh_out_rhand_align - mesh_gt_rhand) ** 2, 1)).mean() * 1000) eval_result['pa_mpvpe_hand'].append((np.sqrt( np.sum((mesh_out_lhand_align - mesh_gt_lhand) ** 2, 1)).mean() * 1000 + np.sqrt( np.sum((mesh_out_rhand_align - mesh_gt_rhand) ** 2, 1)).mean() * 1000) / 2.) # MPVPE from face vertices mesh_gt_face = mesh_gt[smpl_x.face_vertex_idx, :] mesh_out_face = mesh_out[smpl_x.face_vertex_idx, :] mesh_out_face_align = mesh_out_face - np.dot(smpl_x.J_regressor, mesh_out)[smpl_x.J_regressor_idx['neck'], None, :] + np.dot(smpl_x.J_regressor, mesh_gt)[ smpl_x.J_regressor_idx['neck'], None, :] eval_result['mpvpe_face'].append( np.sqrt(np.sum((mesh_out_face_align - mesh_gt_face) ** 2, 1)).mean() * 1000) mesh_out_face_align = rigid_align(mesh_out_face, mesh_gt_face) eval_result['pa_mpvpe_face'].append( np.sqrt(np.sum((mesh_out_face_align - mesh_gt_face) ** 2, 1)).mean() * 1000) if vis: img_path = out['img_path'] rel_img_path = img_path.split('..')[-1] smplx_pred = {} smplx_pred['global_orient'] = out['smplx_root_pose'].reshape(-1,3) smplx_pred['body_pose'] = out['smplx_body_pose'].reshape(-1,3) smplx_pred['left_hand_pose'] = out['smplx_lhand_pose'].reshape(-1,3) smplx_pred['right_hand_pose'] = out['smplx_rhand_pose'].reshape(-1,3) smplx_pred['jaw_pose'] = out['smplx_jaw_pose'].reshape(-1,3) smplx_pred['leye_pose'] = np.zeros((1, 3)) smplx_pred['reye_pose'] = np.zeros((1, 3)) smplx_pred['betas'] = out['smplx_shape'].reshape(-1,10) smplx_pred['expression'] = out['smplx_expr'].reshape(-1,10) smplx_pred['transl'] = out['gt_smplx_transl'].reshape(-1,3) smplx_pred['img_path'] = rel_img_path npz_path = os.path.join(cfg.vis_dir, f'{self.save_idx}.npz') np.savez(npz_path, **smplx_pred) # save img path and error new_line = [self.save_idx, rel_img_path, mpvpe_all, pa_mpvpe_all] # Append the new line to the CSV file writer.writerow(new_line) self.save_idx += 1 # save_obj(out['smplx_mesh_cam'], smpl_x.face, str(cur_sample_idx + n) + '.obj') # save results for the official evaluation codes/server save_name = annot['img_path'].split('/')[-1][:-4] if self.data_split == 'test' and self.test_set == 'test': if self.resolution == (2160, 3840): save_name = save_name.split('_pid')[0] elif self.data_split == 'test' and self.test_set == 'val': if self.resolution == (2160, 3840): save_name = save_name.split('_ann_id')[0] else: save_name = save_name.split('_1280x720')[0] if 'person_idx' in annot: person_idx = annot['person_idx'] else: exist_result_path = glob(osp.join(cfg.result_dir, 'AGORA', save_name + '*')) if len(exist_result_path) == 0: person_idx = 0 else: last_person_idx = max( [int(name.split('personId_')[1].split('.pkl')[0]) for name in exist_result_path]) person_idx = last_person_idx + 1 save_name += '_personId_' + str(person_idx) + '.pkl' joint_proj = out['smplx_joint_proj'] joint_proj[:, 0] = joint_proj[:, 0] / cfg.output_hm_shape[2] * cfg.input_img_shape[1] joint_proj[:, 1] = joint_proj[:, 1] / cfg.output_hm_shape[1] * cfg.input_img_shape[0] joint_proj = np.concatenate((joint_proj, np.ones_like(joint_proj[:, :1])), 1) joint_proj = np.dot(out['bb2img_trans'], joint_proj.transpose(1, 0)).transpose(1, 0) joint_proj[:, 0] = joint_proj[:, 0] / self.resolution[1] * 3840 # restore to original resolution joint_proj[:, 1] = joint_proj[:, 1] / self.resolution[0] * 2160 # restore to original resolution save_dict = {'params': {'transl': out['cam_trans'].reshape(1, -1), 'global_orient': out['smplx_root_pose'].reshape(1, -1), 'body_pose': out['smplx_body_pose'].reshape(1, -1), 'left_hand_pose': out['smplx_lhand_pose'].reshape(1, -1), 'right_hand_pose': out['smplx_rhand_pose'].reshape(1, -1), 'reye_pose': np.zeros((1, 3)), 'leye_pose': np.zeros((1, 3)), 'jaw_pose': out['smplx_jaw_pose'].reshape(1, -1), 'expression': out['smplx_expr'].reshape(1, -1), 'betas': out['smplx_shape'].reshape(1, -1)}, 'joints': joint_proj.reshape(1, -1, 2) } os.makedirs(osp.join(cfg.result_dir, 'predictions'), exist_ok=True) with open(osp.join(cfg.result_dir, 'predictions', save_name), 'wb') as f: pickle.dump(save_dict, f) if getattr(cfg, 'vis', False): file.close() return eval_result def print_eval_result(self, eval_result): print('AGORA test results are dumped at: ' + osp.join(cfg.result_dir, 'predictions')) if self.data_split == 'test' and self.test_set == 'test': # do not print. just submit the results to the official evaluation server return print('======AGORA-val======') print(f'{cfg.vis_dir}') print('PA MPVPE (All): %.2f mm' % np.mean(eval_result['pa_mpvpe_all'])) print('PA MPVPE (L-Hands): %.2f mm' % np.mean(eval_result['pa_mpvpe_l_hand'])) print('PA MPVPE (R-Hands): %.2f mm' % np.mean(eval_result['pa_mpvpe_r_hand'])) print('PA MPVPE (Hands): %.2f mm' % np.mean(eval_result['pa_mpvpe_hand'])) print('PA MPVPE (Face): %.2f mm' % np.mean(eval_result['pa_mpvpe_face'])) print() print('MPVPE (All): %.2f mm' % np.mean(eval_result['mpvpe_all'])) print('MPVPE (L-Hands): %.2f mm' % np.mean(eval_result['mpvpe_l_hand'])) print('MPVPE (R-Hands): %.2f mm' % np.mean(eval_result['mpvpe_r_hand'])) print('MPVPE (Hands): %.2f mm' % np.mean(eval_result['mpvpe_hand'])) print('MPVPE (Face): %.2f mm' % np.mean(eval_result['mpvpe_face'])) print() print(f"{np.mean(eval_result['pa_mpvpe_all'])},{np.mean(eval_result['pa_mpvpe_l_hand'])},{np.mean(eval_result['pa_mpvpe_r_hand'])},{np.mean(eval_result['pa_mpvpe_hand'])},{np.mean(eval_result['pa_mpvpe_face'])}," f"{np.mean(eval_result['mpvpe_all'])},{np.mean(eval_result['mpvpe_l_hand'])},{np.mean(eval_result['mpvpe_r_hand'])},{np.mean(eval_result['mpvpe_hand'])},{np.mean(eval_result['mpvpe_face'])}") print() f = open(os.path.join(cfg.result_dir, 'result.txt'), 'w') f.write(f'AGORA-val dataset: \n') f.write('PA MPVPE (All): %.2f mm\n' % np.mean(eval_result['pa_mpvpe_all'])) f.write('PA MPVPE (L-Hands): %.2f mm' % np.mean(eval_result['pa_mpvpe_l_hand'])) f.write('PA MPVPE (R-Hands): %.2f mm' % np.mean(eval_result['pa_mpvpe_r_hand'])) f.write('PA MPVPE (Hands): %.2f mm\n' % np.mean(eval_result['pa_mpvpe_hand'])) f.write('PA MPVPE (Face): %.2f mm\n' % np.mean(eval_result['pa_mpvpe_face'])) f.write('MPVPE (All): %.2f mm\n' % np.mean(eval_result['mpvpe_all'])) f.write('MPVPE (L-Hands): %.2f mm' % np.mean(eval_result['mpvpe_l_hand'])) f.write('MPVPE (R-Hands): %.2f mm' % np.mean(eval_result['mpvpe_r_hand'])) f.write('MPVPE (Hands): %.2f mm' % np.mean(eval_result['mpvpe_hand'])) f.write('MPVPE (Face): %.2f mm\n' % np.mean(eval_result['mpvpe_face'])) f.write(f"{np.mean(eval_result['pa_mpvpe_all'])},{np.mean(eval_result['pa_mpvpe_l_hand'])},{np.mean(eval_result['pa_mpvpe_r_hand'])},{np.mean(eval_result['pa_mpvpe_hand'])},{np.mean(eval_result['pa_mpvpe_face'])}," f"{np.mean(eval_result['mpvpe_all'])},{np.mean(eval_result['mpvpe_l_hand'])},{np.mean(eval_result['mpvpe_r_hand'])},{np.mean(eval_result['mpvpe_hand'])},{np.mean(eval_result['mpvpe_face'])}") if getattr(cfg, 'eval_on_train', False): import csv csv_file = f'{cfg.root_dir}/output/agora_eval_on_train.csv' exp_id = cfg.exp_name.split('_')[1] new_line = [exp_id,np.mean(eval_result['pa_mpvpe_all']),np.mean(eval_result['pa_mpvpe_l_hand']),np.mean(eval_result['pa_mpvpe_r_hand']),np.mean(eval_result['pa_mpvpe_hand']),np.mean(eval_result['pa_mpvpe_face']), np.mean(eval_result['mpvpe_all']),np.mean(eval_result['mpvpe_l_hand']),np.mean(eval_result['mpvpe_r_hand']),np.mean(eval_result['mpvpe_hand']),np.mean(eval_result['mpvpe_face'])] # Append the new line to the CSV file with open(csv_file, 'a', newline='') as file: writer = csv.writer(file) writer.writerow(new_line)