This module converts the model's output into the format expected by the coco api
| 1195 | |
| 1196 | |
| 1197 | class PostProcess_SMPLX_Multi_Box(nn.Module): |
| 1198 | """ This module converts the model's output into the format expected by the coco api""" |
| 1199 | def __init__( |
| 1200 | self, |
| 1201 | num_select=100, |
| 1202 | nms_iou_threshold=-1, |
| 1203 | num_body_points=17, |
| 1204 | body_model= dict( |
| 1205 | type='smplx', |
| 1206 | keypoint_src='smplx', |
| 1207 | num_expression_coeffs=10, |
| 1208 | num_betas=10, |
| 1209 | gender='neutral', |
| 1210 | keypoint_dst='smplx_137', |
| 1211 | model_path='data/body_models/smplx', |
| 1212 | use_pca=False, |
| 1213 | use_face_contour=True) |
| 1214 | ) -> None: |
| 1215 | super().__init__() |
| 1216 | self.num_select = num_select |
| 1217 | self.nms_iou_threshold = nms_iou_threshold |
| 1218 | self.num_body_points=num_body_points |
| 1219 | |
| 1220 | # -1 for neutral; 0 for male; 1 for femal |
| 1221 | gender_body_model = {} |
| 1222 | gender_body_model[-1] = build_body_model(body_model) |
| 1223 | |
| 1224 | body_model['gender']='male' |
| 1225 | gender_body_model[0] = build_body_model(body_model) |
| 1226 | |
| 1227 | body_model['gender']='female' |
| 1228 | gender_body_model[1] = build_body_model(body_model) |
| 1229 | |
| 1230 | self.body_model = gender_body_model |
| 1231 | |
| 1232 | |
| 1233 | @torch.no_grad() |
| 1234 | def forward(self, outputs, target_sizes, targets, data_batch_nc, not_to_xyxy=False, test=False): |
| 1235 | # import pdb; pdb.set_trace() |
| 1236 | batch_size = outputs['pred_smpl_beta'].shape[0] |
| 1237 | results = [] |
| 1238 | device = outputs['pred_smpl_beta'].device |
| 1239 | for body_model in self.body_model.values(): |
| 1240 | body_model.to(device) |
| 1241 | # test with instance num |
| 1242 | # num_select=data_batch_nc['joint_img'][0].shape[0] |
| 1243 | num_select = self.num_select |
| 1244 | out_logits, out_bbox= outputs['pred_logits'], outputs['pred_boxes'] |
| 1245 | |
| 1246 | out_smpl_pose, out_smpl_beta, out_smpl_expr, out_smpl_cam, out_smpl_kp3d, out_smpl_verts = \ |
| 1247 | outputs['pred_smpl_fullpose'], outputs['pred_smpl_beta'], outputs['pred_smpl_expr'], \ |
| 1248 | outputs['pred_smpl_cam'], outputs['pred_smpl_kp3d'], outputs['pred_smpl_verts'] |
| 1249 | |
| 1250 | out_smpl_kp2d = [] |
| 1251 | |
| 1252 | for bs in range(batch_size): |
| 1253 | out_kp3d_i = out_smpl_kp3d[bs] |
| 1254 | out_cam_i = out_smpl_cam[bs] |
no outgoing calls
no test coverage detected