import os import os.path as osp import numpy as np import torch import cv2 import json import copy from pycocotools.coco import COCO from config import cfg from utils.human_models import smpl_x from utils.preprocessing import load_img, process_bbox, augmentation, process_db_coord, process_human_model_output, \ get_fitting_error_3D from utils.transforms import world2cam, cam2pixel, rigid_align from humandata import HumanDataset import pickle from body_measurements import BodyMeasurements import smplx from test_submission_format import test_submission_file_format def point_error(x, y, align=True): """ Ref: https://github.com/muelea/shapy/blob/master/regressor/hbw_evaluation/evaluate_hbw.py#LL44C1-L58C31 """ t = 0.0 if align: t = x.mean(0, keepdims=True) - y.mean(0, keepdims=True) x_hat = x - t error = np.sqrt(np.power(x_hat - y, 2).sum(axis=-1)) return error.mean().item() class SHAPY(HumanDataset): def __init__(self, transform, data_split): super(SHAPY, self).__init__(transform, data_split) self.eval_split = getattr(cfg, 'shapy_eval_split') if self.data_split == 'train': raise NotImplementedError('Shapy train not implemented yet. Need to consider invalid parameters') if self.data_split == 'test' and self.eval_split == 'test': filename = getattr(cfg, 'filename', 'shapy_test_230512_1631.npz') elif self.data_split == 'test' and self.eval_split == 'val': filename = getattr(cfg, 'filename', 'shapy_val_230512_705.npz') else: raise ValueError(f'Undefined. data split: {self.data_split}; eval_split: {self.test_set}') self.img_dir = osp.join(cfg.data_dir, 'SHAPY') self.annot_path = osp.join(cfg.data_dir, 'preprocessed_datasets', filename) self.v_shape_load_dir = osp.join(cfg.data_dir, 'SHAPY', 'HBW', 'smplx', 'val') self.img_shape = None # variable img_shape self.cam_param = {} # load data self.datalist = self.load_data( train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1)) ### SHAPY utils ### ref: https://github.com/muelea/shapy/blob/master/regressor/hbw_evaluation/evaluate_hbw.py#L28 # load body model # ref: common/utils/human_models.py self.layer_arg = {'create_global_orient': False, 'create_body_pose': False, 'create_left_hand_pose': False, 'create_right_hand_pose': False, 'create_jaw_pose': False, 'create_leye_pose': False, 'create_reye_pose': False, 'create_betas': False, 'create_expression': False, 'create_transl': False} self.smplx_layer = smplx.create(cfg.human_model_path, 'smplx', gender='NEUTRAL', use_pca=False, use_face_contour=True, flat_hand_mean=True, # critical! **self.layer_arg).cuda() # self.smplx_layer = copy.deepcopy(smpl_x.layer['neutral']).cuda() self.faces_tensor_smplx = self.smplx_layer.faces_tensor.detach().cpu().numpy() # load files to compute P2P-20K Error point_reg = osp.join(cfg.data_dir, 'SHAPY', 'utility_files', 'evaluation', 'eval_point_set', 'HD_SMPLX_from_SMPL.pkl') with open(point_reg, 'rb') as f: self.point_regressor = pickle.load(f) # load files to compute Measurements Error body_measurement_folder = osp.join(cfg.data_dir, 'SHAPY', 'utility_files', 'measurements') meas_def_path = osp.join(body_measurement_folder, 'measurement_defitions.yaml') meas_verts_path_gt = osp.join(body_measurement_folder, 'smplx_measurements.yaml') self.body_measurements = BodyMeasurements( {'meas_definition_path': meas_def_path, 'meas_vertices_path': meas_verts_path_gt}, ).to('cuda') self.v_shaped_gt = {} # to save preditions self.images_names = [] self.v_shaped = [] def evaluate(self, outs, cur_sample_idx): annots = self.datalist sample_num = len(outs) eval_result = {'v2v_t_errors': [], 'point_t_errors': [], 'height': [], 'chest': [], 'waist': [], 'hips': [], 'mass': []} for n in range(sample_num): annot = annots[cur_sample_idx + n] out = outs[n] betas_fit = out['smplx_shape'] img_path = out['img_path'] # compute v_shaped betas_fit = torch.tensor(betas_fit.reshape(-1, 10)).cuda() output = self.smplx_layer( betas=betas_fit, body_pose=torch.zeros((1, 63)).to(betas_fit.device), global_orient=torch.zeros((1, 3)).to(betas_fit.device), right_hand_pose=torch.zeros((1, 45)).to(betas_fit.device), left_hand_pose=torch.zeros((1, 45)).to(betas_fit.device), jaw_pose=torch.zeros((1, 3)).to(betas_fit.device), leye_pose=torch.zeros((1, 3)).to(betas_fit.device), reye_pose=torch.zeros((1, 3)).to(betas_fit.device), expression=torch.zeros((1, 10)).to(betas_fit.device), return_verts=True ) v_shaped_fit = output.vertices.detach().cpu().numpy().squeeze() image_name = '/'.join(img_path.split('/')[-4:]) self.images_names.append(image_name) self.v_shaped.append(v_shaped_fit) if self.eval_split == 'val': # load gt vertices subject = img_path.split('/')[-3] subject_id_npy = subject.split('_')[0] + '.npy' v_shaped_gt_path = osp.join(self.v_shape_load_dir, subject_id_npy) if v_shaped_gt_path not in self.v_shaped_gt: v_shaped_gt = np.load(v_shaped_gt_path) self.v_shaped_gt[v_shaped_gt_path] = v_shaped_gt else: v_shaped_gt = self.v_shaped_gt[v_shaped_gt_path] # compute vertex-to-vertex error (SMPL-X only) # ref: https://github.com/muelea/shapy/blob/master/regressor/hbw_evaluation/evaluate_hbw.py#LL142C1-L171C48 v2v_error = point_error(v_shaped_fit, v_shaped_gt, align=True) eval_result['v2v_t_errors'].append(v2v_error) # compute P2P-20k error points_gt = self.point_regressor.dot(v_shaped_gt) points_fit = self.point_regressor.dot(v_shaped_fit) p2p_error = point_error(points_gt, points_fit, align=True) eval_result['point_t_errors'].append(p2p_error) # compute height/chest/waist/hip error shaped_triangles_gt = v_shaped_gt[self.faces_tensor_smplx] shaped_triangles_gt = torch.from_numpy(shaped_triangles_gt).unsqueeze(0).to('cuda') measurements_gt = self.body_measurements(shaped_triangles_gt)['measurements'] shaped_triangles_fit = v_shaped_fit[self.faces_tensor_smplx] shaped_triangles_fit = torch.from_numpy(shaped_triangles_fit).unsqueeze(0).to('cuda') measurements_fit = self.body_measurements(shaped_triangles_fit)['measurements'] for k in ['height', 'chest', 'waist', 'hips', 'mass']: error = abs(measurements_gt[k]['tensor'].item() - measurements_fit[k]['tensor'].item()) eval_result[k].append(error) return eval_result def print_eval_result(self, eval_result): # print('SHAPY results are dumped at: ' + osp.join(cfg.result_dir, 'predictions')) if self.data_split == 'test' and self.eval_split == 'test': # do not print. just submit the results to the official evaluation server # save predictions in the format of HBW challenge # ref: https://github.com/muelea/shapy/blob/master/regressor/hbw_evaluation/README_HBW_EVAL.md#hbw-challenge save_dir = osp.join(cfg.result_dir, 'predictions') os.makedirs(save_dir, exist_ok=True) save_name = osp.join(save_dir, 'hbw_prediction') images_names = np.array(self.images_names).reshape(1631, ) v_shaped = np.array(self.v_shaped).reshape(1631, 10475, 3) np.savez(save_name, image_name=images_names, v_shaped=v_shaped) print('predictions saved at: ' + save_name + '.npz') # run format test test_submission_file_format(save_name + '.npz') return v2v_t_errors = np.mean(eval_result['v2v_t_errors']) * 1000 point_t_errors = np.mean(eval_result['point_t_errors']) * 1000 chest = np.mean(eval_result['chest']) * 1000 waist = np.mean(eval_result['waist']) * 1000 hips = np.mean(eval_result['hips']) * 1000 height = np.mean(eval_result['height']) * 1000 mass = np.mean(eval_result['mass']) print('======SHAPY-val======') print('Height Error: %.2f mm' % height) print('Chest Error: %.2f mm' % chest) print('Waist Error: %.2f mm' % waist) print('Hips Error: %.2f mm' % hips) print('P2P-20k Error: %.2f mm' % point_t_errors) print('V2V Error: %.2f mm' % v2v_t_errors) print('Mass Error: %.2f kg' % mass) f = open(os.path.join(cfg.result_dir, 'result.txt'), 'w') f.write(f'SHAPY-val dataset: \n') f.write('Height Error: %.2f mm\n' % height) f.write('Chest Error: %.2f mm' % chest) f.write('Waist Error: %.2f mm\n' % waist) f.write('Hips Error: %.2f mm\n' % hips) f.write('P2P-20k Error: %.2f mm' % point_t_errors) f.write('V2V Error: %.2f mm\n' % v2v_t_errors) f.write('Mass Error: %.2f kg\n' % mass) f.close()