import os import os.path as osp from glob import glob import numpy as np from config.config import cfg import copy import json import cv2 import torch from pycocotools.coco import COCO from util.human_models import smpl_x from util.preprocessing import load_img, process_bbox, load_ply from util.transforms import rigid_align, rigid_align_batch from humandata import HumanDataset import csv class EHF(HumanDataset): def __init__(self, transform, data_split): super(EHF, self).__init__(transform, data_split) self.transform = transform self.data_split = data_split self.save_idx = 0 # self.cam_param = {'R': [-2.98747896, 0.01172457, -0.05704687]} # self.cam_param['R'], _ = cv2.Rodrigues(np.array(self.cam_param['R'])) self.cam_param = {} self.img_dir = 'data/data_weichen/ehf' self.img_shape = [1200, 1600] self.annot_path = 'data_tmp/multihuman_data/ehf_val_230908_100.npz' self.annot_path_cache = 'data_tmp/cache/ehf_val_cache_230908_100.npz' if self.use_cache and osp.isfile(self.annot_path_cache): print(f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}') self.datalist = self.load_cache(self.annot_path_cache) else: if self.use_cache: print(f'[{self.__class__.__name__}] Cache not found, generating cache...') self.datalist = self.load_data( train_sample_interval=getattr(cfg, f'{self.__class__.__name__}_train_sample_interval', 1)) if self.use_cache: self.save_cache(self.annot_path_cache, self.datalist) 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': [], 'pa_mpjpe_body': [], 'pa_mpjpe_l_hand': [], 'pa_mpjpe_r_hand': [], 'pa_mpjpe_hand': [] } csv_file = f'{cfg.result_dir}/ehf_smplx_error.csv' file = open(csv_file, 'a', newline='') for n in range(sample_num): annot = annots[cur_sample_idx + n] ann_id = annot['img_path'].split('/')[-1].split('_')[0] out = outs[n] ann_idx = out['gt_ann_idx'] img_path = [] for ann_id in ann_idx: img_path.append(annots[ann_id]['img_path']) eval_result['img_path'] = img_path eval_result['ann_idx'] = ann_idx # MPVPE from all vertices np.dot(self.cam_param['R'], out['smplx_mesh_cam_target'].transpose(0,2,1)).transpose(1,2,0) # mesh_gt = np.dot( # self.cam_param['R'], # out['smplx_mesh_cam_target'].transpose(0,2,1) # ).transpose(1,2,0) mesh_gt = out['smplx_mesh_cam_target'] mesh_out = out['smplx_mesh_cam'] # mesh_gt_align = rigid_align(mesh_gt, mesh_out) # print(mesh_out.shape) mesh_out_align = rigid_align_batch(mesh_out, mesh_gt) eval_result['pa_mpvpe_all'].append( np.sqrt(np.sum( (mesh_out_align - mesh_gt)**2, -1)).mean() * 1000) mesh_out_align = mesh_out - np.dot( smpl_x.J_regressor, mesh_out).transpose(1,0,2)[:, smpl_x.J_regressor_idx['pelvis'], None, :] + np.dot( smpl_x.J_regressor, mesh_gt).transpose(1,0,2)[:, smpl_x.J_regressor_idx['pelvis'], None, :] eval_result['mpvpe_all'].append( np.sqrt(np.sum( (mesh_out_align - mesh_gt)**2, -1)).mean() * 1000) # 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_out_lhand_align = rigid_align_batch(mesh_out_lhand, mesh_gt_lhand) 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_rhand_align = rigid_align_batch(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.) mesh_out_lhand_align = mesh_out_lhand - np.dot( smpl_x.J_regressor, mesh_out).transpose(1,0,2)[:, smpl_x.J_regressor_idx['lwrist'], None, :] + np.dot( smpl_x.J_regressor, mesh_gt).transpose(1,0,2)[:, smpl_x.J_regressor_idx['lwrist'], None, :] mesh_out_rhand_align = mesh_out_rhand - np.dot( smpl_x.J_regressor, mesh_out).transpose(1,0,2)[:, smpl_x.J_regressor_idx['rwrist'], None, :] + np.dot( smpl_x.J_regressor, mesh_gt).transpose(1,0,2)[:, 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.) # 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 = rigid_align_batch(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) mesh_out_face_align = mesh_out_face - np.dot( smpl_x.J_regressor, mesh_out).transpose(1,0,2)[:, smpl_x.J_regressor_idx['neck'], None, :] + np.dot( smpl_x.J_regressor, mesh_gt).transpose(1,0,2)[:, 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) # MPJPE from body joints joint_gt_body = np.dot(smpl_x.j14_regressor, mesh_gt).transpose(1,0,2) joint_out_body = np.dot(smpl_x.j14_regressor, mesh_out).transpose(1,0,2) joint_out_body_align = rigid_align_batch(joint_out_body, joint_gt_body) eval_result['pa_mpjpe_body'].append( np.sqrt(np.sum( (joint_out_body_align - joint_gt_body)**2, -1)).mean() * 1000) # MPJPE from hand joints joint_gt_lhand = np.dot(smpl_x.orig_hand_regressor['left'], mesh_gt).transpose(1,0,2) joint_out_lhand = np.dot(smpl_x.orig_hand_regressor['left'], mesh_out).transpose(1,0,2) joint_out_lhand_align = rigid_align_batch(joint_out_lhand, joint_gt_lhand) joint_gt_rhand = np.dot(smpl_x.orig_hand_regressor['right'], mesh_gt).transpose(1,0,2) joint_out_rhand = np.dot(smpl_x.orig_hand_regressor['right'], mesh_out).transpose(1,0,2) joint_out_rhand_align = rigid_align_batch(joint_out_rhand, joint_gt_rhand) eval_result['pa_mpjpe_l_hand'].append( np.sqrt(np.sum( (joint_out_lhand_align - joint_gt_lhand)**2, -1)).mean() * 1000) eval_result['pa_mpjpe_r_hand'].append( np.sqrt(np.sum( (joint_out_rhand_align - joint_gt_rhand)**2, 1)).mean() * 1000) eval_result['pa_mpjpe_hand'].append( (np.sqrt(np.sum( (joint_out_lhand_align - joint_gt_lhand)**2, -1)).mean() * 1000 + np.sqrt(np.sum( (joint_out_rhand_align - joint_gt_rhand)**2, -1)).mean() * 1000) / 2.) save_error=True if save_error: writer = csv.writer(file) new_line = [ann_idx[n],img_path[n], eval_result['mpvpe_all'][-1], eval_result['pa_mpvpe_all'][-1]] writer.writerow(new_line) self.save_idx += 1 # vis = cfg.vis for k,v in eval_result.items(): if k != 'img_path' and k != 'ann_idx': if len(v)>1: eval_result[k] = np.concatenate(v,axis=0) else: eval_result[k] = np.array(v) return eval_result def print_eval_result(self, eval_result): print('======EHF======') 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('PA MPJPE (Body): %.2f mm' % np.mean(eval_result['pa_mpjpe_body'])) print('PA MPJPE (L-Hands): %.2f mm' % np.mean(eval_result['pa_mpjpe_l_hand'])) print('PA MPJPE (R-Hands): %.2f mm' % np.mean(eval_result['pa_mpjpe_r_hand'])) print('PA MPJPE (Hands): %.2f mm' % np.mean(eval_result['pa_mpjpe_hand'])) out_file = osp.join(cfg.result_dir,'ehf_test.txt') if os.path.exists(out_file): f = open(out_file, 'a+') else: f = open(out_file, 'w', encoding="utf-8") f.write('\n') f.write(f'{cfg.exp_name}\n') f.write(f'EHF 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\n' % np.mean(eval_result['pa_mpvpe_l_hand'])) f.write('PA MPVPE (R-Hands): %.2f mm\n' % 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\n' % np.mean(eval_result['mpvpe_l_hand'])) f.write('MPVPE (R-Hands): %.2f mm\n' % np.mean(eval_result['mpvpe_r_hand'])) f.write('MPVPE (Hands): %.2f mm\n' % np.mean(eval_result['mpvpe_hand'])) f.write('MPVPE (Face): %.2f mm\n' % np.mean(eval_result['mpvpe_face'])) f.write('PA MPJPE (Body): %.2f mm\n' % np.mean(eval_result['pa_mpjpe_body'])) f.write('PA MPJPE (L-Hands): %.2f mm\n' % np.mean(eval_result['pa_mpjpe_l_hand'])) f.write('PA MPJPE (R-Hands): %.2f mm\n' % np.mean(eval_result['pa_mpjpe_r_hand'])) f.write('PA MPJPE (Hands): %.2f mm\n' % np.mean(eval_result['pa_mpjpe_hand'])) f.close()