This module converts the model's output into the format expected by the coco api
| 1512 | |
| 1513 | |
| 1514 | class PostProcess_SMPLX_Multi_Infer_Box(nn.Module): |
| 1515 | """ This module converts the model's output into the format expected by the coco api""" |
| 1516 | def __init__( |
| 1517 | self, |
| 1518 | num_select=100, |
| 1519 | nms_iou_threshold=-1, |
| 1520 | num_body_points=17, |
| 1521 | body_model= dict( |
| 1522 | type='smplx', |
| 1523 | keypoint_src='smplx', |
| 1524 | num_expression_coeffs=10, |
| 1525 | num_betas=10, |
| 1526 | gender='neutral', |
| 1527 | keypoint_dst='smplx_137', |
| 1528 | model_path='data/body_models/smplx', |
| 1529 | use_pca=False, |
| 1530 | use_face_contour=True) |
| 1531 | ) -> None: |
| 1532 | super().__init__() |
| 1533 | self.num_select = num_select |
| 1534 | self.nms_iou_threshold = nms_iou_threshold |
| 1535 | self.num_body_points=num_body_points |
| 1536 | |
| 1537 | # -1 for neutral; 0 for male; 1 for femal |
| 1538 | gender_body_model = {} |
| 1539 | gender_body_model[-1] = build_body_model(body_model) |
| 1540 | |
| 1541 | body_model['gender']='male' |
| 1542 | gender_body_model[0] = build_body_model(body_model) |
| 1543 | |
| 1544 | body_model['gender']='female' |
| 1545 | gender_body_model[1] = build_body_model(body_model) |
| 1546 | |
| 1547 | self.body_model = gender_body_model |
| 1548 | |
| 1549 | @torch.no_grad() |
| 1550 | def forward(self, outputs, target_sizes, targets, data_batch_nc, image_shape= None, not_to_xyxy=False, test=False): |
| 1551 | """ |
| 1552 | image_shape(target_sizes): input image shape |
| 1553 | |
| 1554 | """ |
| 1555 | |
| 1556 | batch_size = outputs['pred_smpl_beta'].shape[0] |
| 1557 | results = [] |
| 1558 | device = outputs['pred_smpl_beta'].device |
| 1559 | |
| 1560 | num_select = self.num_select |
| 1561 | out_logits, out_bbox= outputs['pred_logits'], outputs['pred_boxes'] |
| 1562 | |
| 1563 | out_body_bbox, out_lhand_bbox, out_rhand_bbox, out_face_bbox = \ |
| 1564 | outputs['pred_boxes'], outputs['pred_lhand_boxes'], \ |
| 1565 | outputs['pred_rhand_boxes'], outputs['pred_face_boxes'] |
| 1566 | |
| 1567 | out_smpl_pose, out_smpl_beta, out_smpl_expr, out_smpl_cam, out_smpl_kp3d, out_smpl_verts = \ |
| 1568 | outputs['pred_smpl_fullpose'], outputs['pred_smpl_beta'], outputs['pred_smpl_expr'], \ |
| 1569 | outputs['pred_smpl_cam'], outputs['pred_smpl_kp3d'], outputs['pred_smpl_verts'] |
| 1570 | |
| 1571 | out_smpl_kp2d = [] |
no outgoing calls
no test coverage detected