This module converts the model's output into the format expected by the coco api
| 380 | |
| 381 | |
| 382 | class PostProcess_SMPLX(nn.Module): |
| 383 | """ This module converts the model's output into the format expected by the coco api""" |
| 384 | def __init__( |
| 385 | self, |
| 386 | num_select=100, |
| 387 | nms_iou_threshold=-1, |
| 388 | num_body_points=17, |
| 389 | body_model= dict( |
| 390 | type='smplx', |
| 391 | keypoint_src='smplx', |
| 392 | num_expression_coeffs=10, |
| 393 | keypoint_dst='smplx_137', |
| 394 | model_path='data/body_models/smplx', |
| 395 | use_pca=False, |
| 396 | use_face_contour=True) |
| 397 | ) -> None: |
| 398 | super().__init__() |
| 399 | self.num_select = num_select |
| 400 | self.nms_iou_threshold = nms_iou_threshold |
| 401 | self.num_body_points=num_body_points |
| 402 | self.body_model = build_body_model(body_model) |
| 403 | |
| 404 | @torch.no_grad() |
| 405 | def forward(self, outputs, target_sizes, targets, data_batch_nc, not_to_xyxy=False, test=False): |
| 406 | # import pdb; pdb.set_trace() |
| 407 | num_select = self.num_select |
| 408 | |
| 409 | out_logits, out_bbox, out_keypoints= \ |
| 410 | outputs['pred_logits'], outputs['pred_boxes'], \ |
| 411 | outputs['pred_keypoints'] |
| 412 | |
| 413 | out_smpl_pose, out_smpl_beta, out_smpl_expr, out_smpl_cam, out_smpl_kp3d, out_smpl_verts = \ |
| 414 | outputs['pred_smpl_fullpose'], outputs['pred_smpl_beta'], outputs['pred_smpl_expr'], \ |
| 415 | outputs['pred_smpl_cam'], outputs['pred_smpl_kp3d'], outputs['pred_smpl_verts'] |
| 416 | |
| 417 | assert len(out_logits) == len(target_sizes) |
| 418 | assert target_sizes.shape[1] == 2 |
| 419 | prob = out_logits.sigmoid() |
| 420 | topk_values, topk_indexes = \ |
| 421 | torch.topk(prob.view(out_logits.shape[0], -1), num_select, dim=1) |
| 422 | scores = topk_values |
| 423 | # bbox |
| 424 | topk_boxes = topk_indexes // out_logits.shape[2] |
| 425 | labels = topk_indexes % out_logits.shape[2] |
| 426 | if not_to_xyxy: |
| 427 | boxes = out_bbox |
| 428 | else: |
| 429 | boxes = box_ops.box_cxcywh_to_xyxy(out_bbox) |
| 430 | |
| 431 | if test: |
| 432 | assert not not_to_xyxy |
| 433 | boxes[:,:,2:] = boxes[:,:,2:] - boxes[:,:,:2] |
| 434 | boxes_norm = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1,1,4)) |
| 435 | target_sizes = target_sizes.type_as(boxes) |
| 436 | # from relative [0, 1] to absolute [0, height] coordinates |
| 437 | img_h, img_w = target_sizes.unbind(1) |
| 438 | scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1) |
| 439 | boxes = boxes_norm * scale_fct[:, None, :] |