| 55 | return result |
| 56 | |
| 57 | def inference(self, outs): |
| 58 | img_paths = self.img_paths |
| 59 | output = {} |
| 60 | |
| 61 | for out in outs: |
| 62 | ann_idx = out['image_idx'] |
| 63 | scores = out['scores'].clone().cpu().numpy() |
| 64 | img_shape = out['img_shape'].cpu().numpy()[::-1] # w, h |
| 65 | img = cv2.imread(img_paths[ann_idx]) # h, w |
| 66 | joint_proj = out['smplx_joint_proj'].clone().cpu().numpy() |
| 67 | scale = img.shape[1]/img_shape[0] |
| 68 | joint_proj *= scale |
| 69 | |
| 70 | for i, score in enumerate(scores): |
| 71 | if score < self.score_threshold: |
| 72 | break |
| 73 | save_name = img_paths[ann_idx].split('/')[-1][:-4] |
| 74 | if self.resolution == (2160, 3840): |
| 75 | save_name = save_name.split('_ann_id')[0] |
| 76 | else: |
| 77 | save_name = save_name.split('_1280x720')[0] |
| 78 | |
| 79 | save_dict = { |
| 80 | 'params': { |
| 81 | 'transl': out['cam_trans'][i].reshape(1, -1).cpu().numpy(), |
| 82 | 'global_orient': out['smplx_root_pose'][i].reshape(1, -1).cpu().numpy(), |
| 83 | 'body_pose': out['smplx_body_pose'][i].reshape(1, -1).cpu().numpy(), |
| 84 | 'left_hand_pose': out['smplx_lhand_pose'][i].reshape(1, -1).cpu().numpy(), |
| 85 | 'right_hand_pose': out['smplx_rhand_pose'][i].reshape(1, -1).cpu().numpy(), |
| 86 | 'reye_pose': np.zeros((1, 3)), |
| 87 | 'leye_pose': np.zeros((1, 3)), |
| 88 | 'jaw_pose': out['smplx_jaw_pose'][i].reshape(1, -1).cpu().numpy(), |
| 89 | 'expression': out['smplx_expr'][i].reshape(1, -1).cpu().numpy(), |
| 90 | 'betas': out['smplx_shape'][i].reshape(1, -1).cpu().numpy()}, |
| 91 | 'joints': joint_proj[i].reshape(1, -1, 2)[0,:24]} |
| 92 | |
| 93 | # save |
| 94 | exist_result_path = glob(osp.join(self.out_path, 'predictions', save_name + '*')) |
| 95 | if len(exist_result_path) == 0: |
| 96 | person_idx = 0 |
| 97 | else: |
| 98 | last_person_idx = max([ |
| 99 | int(name.split('personId_')[1].split('.pkl')[0]) |
| 100 | for name in exist_result_path |
| 101 | ]) |
| 102 | person_idx = last_person_idx + 1 |
| 103 | save_name += '_personId_' + str(person_idx) + '.pkl' |
| 104 | os.makedirs(osp.join(self.out_path, 'predictions'), exist_ok=True) |
| 105 | with open(osp.join(self.out_path, 'predictions', save_name),'wb') as f: |
| 106 | pickle.dump(save_dict, f) |
| 107 | return output |
| 108 | |