| 7 | |
| 8 | |
| 9 | class HungarianMatcher(nn.Module): |
| 10 | def __init__(self, |
| 11 | cost_class: float = 1, |
| 12 | cost_bbox: float = 1, |
| 13 | cost_giou: float = 1, |
| 14 | focal_alpha=0.25, |
| 15 | cost_keypoints=1.0, |
| 16 | cost_kpvis=0.1, |
| 17 | cost_oks=0.01, |
| 18 | num_body_points=17): |
| 19 | super().__init__() |
| 20 | self.cost_class = cost_class |
| 21 | self.cost_bbox = cost_bbox |
| 22 | self.cost_giou = cost_giou |
| 23 | assert cost_class != 0 or cost_bbox != 0 or cost_giou != 0, 'all costs cant be 0' |
| 24 | self.cost_keypoints = cost_keypoints |
| 25 | self.cost_kpvis = cost_kpvis |
| 26 | self.cost_oks = cost_oks |
| 27 | self.focal_alpha = focal_alpha |
| 28 | self.num_body_points = num_body_points |
| 29 | if num_body_points == 17: |
| 30 | self.sigmas = np.array([ |
| 31 | .26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, |
| 32 | 1.07, .87, .87, .89, .89 |
| 33 | ], |
| 34 | dtype=np.float32) / 10.0 |
| 35 | |
| 36 | elif num_body_points == 14: |
| 37 | self.sigmas = np.array([ |
| 38 | .79, .79, .72, .72, .62, .62, 1.07, 1.07, .87, .87, .89, .89, |
| 39 | .79, .79 |
| 40 | ]) / 10.0 |
| 41 | else: |
| 42 | raise ValueError(f'Unsupported keypoints number {num_body_points}') |
| 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 |