| 43 | |
| 44 | @torch.no_grad() |
| 45 | def forward(self, |
| 46 | outputs, |
| 47 | target_sizes, |
| 48 | targets, |
| 49 | data_batch_nc, |
| 50 | device, |
| 51 | not_to_xyxy=False, |
| 52 | test=False): |
| 53 | # import pdb; pdb.set_trace() |
| 54 | num_select = self.num_select |
| 55 | self.body_model.to(device) |
| 56 | |
| 57 | out_logits, out_bbox, out_keypoints= \ |
| 58 | outputs['pred_logits'], outputs['pred_boxes'], \ |
| 59 | outputs['pred_keypoints'] |
| 60 | |
| 61 | out_smpl_pose, out_smpl_beta, out_smpl_cam, out_smpl_kp3d = \ |
| 62 | outputs['pred_smpl_pose'], outputs['pred_smpl_beta'], \ |
| 63 | outputs['pred_smpl_cam'], outputs['pred_smpl_kp3d'] |
| 64 | |
| 65 | assert len(out_logits) == len(target_sizes) |
| 66 | assert target_sizes.shape[1] == 2 |
| 67 | prob = out_logits.sigmoid() |
| 68 | topk_values, topk_indexes = \ |
| 69 | torch.topk(prob.view(out_logits.shape[0], -1), num_select, dim=1) |
| 70 | scores = topk_values |
| 71 | # bbox |
| 72 | topk_boxes = topk_indexes // out_logits.shape[2] |
| 73 | labels = topk_indexes % out_logits.shape[2] |
| 74 | if not_to_xyxy: |
| 75 | boxes = out_bbox |
| 76 | else: |
| 77 | boxes = box_ops.box_cxcywh_to_xyxy(out_bbox) |
| 78 | |
| 79 | if test: |
| 80 | assert not not_to_xyxy |
| 81 | boxes[:, :, 2:] = boxes[:, :, 2:] - boxes[:, :, :2] |
| 82 | boxes_norm = torch.gather(boxes, 1, |
| 83 | topk_boxes.unsqueeze(-1).repeat(1, 1, 4)) |
| 84 | target_sizes = target_sizes.type_as(boxes) |
| 85 | # from relative [0, 1] to absolute [0, height] coordinates |
| 86 | img_h, img_w = target_sizes.unbind(1) |
| 87 | scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1) |
| 88 | boxes = boxes_norm * scale_fct[:, None, :] |
| 89 | |
| 90 | # keypoints |
| 91 | topk_keypoints = topk_indexes // out_logits.shape[2] |
| 92 | labels = topk_indexes % out_logits.shape[2] |
| 93 | keypoints = torch.gather( |
| 94 | out_keypoints, 1, |
| 95 | topk_keypoints.unsqueeze(-1).repeat(1, 1, |
| 96 | self.num_body_points * 3)) |
| 97 | |
| 98 | Z_pred = keypoints[:, :, :(self.num_body_points * 2)] |
| 99 | V_pred = keypoints[:, :, (self.num_body_points * 2):] |
| 100 | img_h, img_w = target_sizes.unbind(1) |
| 101 | Z_pred = Z_pred * torch.stack([img_w, img_h], dim=1).repeat( |
| 102 | 1, self.num_body_points)[:, None, :] |