(self, outputs, targets, data_batch=None)
| 43 | |
| 44 | @torch.no_grad() |
| 45 | def forward(self, outputs, targets, data_batch=None): |
| 46 | bs, num_queries = outputs['pred_logits'].shape[:2] |
| 47 | out_prob = outputs['pred_logits'].flatten(0, 1).sigmoid() |
| 48 | out_bbox = outputs['pred_boxes'].flatten(0, 1) |
| 49 | |
| 50 | out_keypoints = outputs['pred_keypoints'].flatten(0, 1) |
| 51 | |
| 52 | # Also concat the target labels and boxes |
| 53 | tgt_ids = torch.cat([v['labels'] for v in targets]) |
| 54 | tgt_bbox = torch.cat([v['boxes'] for v in targets]) |
| 55 | tgt_keypoints = torch.cat([v['keypoints'] for v in targets]) |
| 56 | tgt_area = torch.cat([v['area'] for v in targets]) |
| 57 | # Compute the classification cost. |
| 58 | alpha = self.focal_alpha |
| 59 | gamma = 2.0 |
| 60 | neg_cost_class = (1 - alpha) * (out_prob** |
| 61 | gamma) * (-(1 - out_prob + 1e-8).log()) |
| 62 | pos_cost_class = alpha * ( |
| 63 | (1 - out_prob)**gamma) * (-(out_prob + 1e-8).log()) |
| 64 | cost_class = pos_cost_class[:, tgt_ids] - neg_cost_class[:, tgt_ids] |
| 65 | |
| 66 | # Compute the L1 cost between boxes |
| 67 | cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) |
| 68 | |
| 69 | # Compute the giou cost betwen boxes |
| 70 | |
| 71 | cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), |
| 72 | box_cxcywh_to_xyxy(tgt_bbox), data_batch) |
| 73 | |
| 74 | # compute the keypoint costs |
| 75 | Z_pred = out_keypoints[:, 0:(self.num_body_points * 2)] |
| 76 | V_pred = out_keypoints[:, (self.num_body_points * 2):] |
| 77 | Z_gt = tgt_keypoints[:, 0:(self.num_body_points * 2)] |
| 78 | V_gt: torch.Tensor = tgt_keypoints[:, (self.num_body_points * 2):] |
| 79 | if Z_pred.sum() > 0: |
| 80 | sigmas = Z_pred.new_tensor(self.sigmas) |
| 81 | variances = (sigmas * 2)**2 |
| 82 | kpt_preds = Z_pred.reshape(-1, Z_pred.size(-1) // 2, 2) |
| 83 | kpt_gts = Z_gt.reshape(-1, Z_gt.size(-1) // 2, 2) |
| 84 | squared_distance = (kpt_preds[:, None, :, 0] - kpt_gts[None, :, :, 0]) ** 2 + \ |
| 85 | (kpt_preds[:, None, :, 1] - kpt_gts[None, :, :, 1]) ** 2 |
| 86 | squared_distance0 = squared_distance / (tgt_area[:, None] * |
| 87 | variances[None, :] * 2) |
| 88 | squared_distance1 = torch.exp(-squared_distance0) |
| 89 | squared_distance1 = squared_distance1 * V_gt |
| 90 | oks = squared_distance1.sum(dim=-1) / (V_gt.sum(dim=-1) + 1e-6) |
| 91 | oks = oks.clamp(min=1e-6) |
| 92 | cost_oks = 1 - oks |
| 93 | # import pdb; pdb.set_trace() |
| 94 | cost_keypoints = torch.abs(Z_pred[:, None, :] - Z_gt[None]) |
| 95 | cost_keypoints = cost_keypoints * V_gt.repeat_interleave( |
| 96 | 2, dim=1)[None] |
| 97 | cost_keypoints = cost_keypoints.sum(-1) |
| 98 | cost_bbox = torch.zeros_like(cost_keypoints) |
| 99 | cost_giou = torch.zeros_like( |
| 100 | cost_keypoints) # [bs*query, instance_num] |
| 101 | C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou + self.cost_keypoints * cost_keypoints + self.cost_oks * cost_oks |
| 102 | C = C.view(bs, num_queries, -1).cpu() |
nothing calls this directly
no test coverage detected