| 155 | |
| 156 | |
| 157 | class HungarianMatcherBox(nn.Module): |
| 158 | def __init__(self, |
| 159 | cost_class: float = 1, |
| 160 | cost_bbox: float = 1, |
| 161 | cost_giou: float = 1, |
| 162 | focal_alpha=0.25, |
| 163 | cost_keypoints=1.0, |
| 164 | cost_kpvis=0.1, |
| 165 | cost_oks=0.01, |
| 166 | num_body_points=17): |
| 167 | super().__init__() |
| 168 | self.cost_class = cost_class |
| 169 | self.cost_bbox = cost_bbox |
| 170 | self.cost_giou = cost_giou |
| 171 | assert cost_class != 0 or cost_bbox != 0 or cost_giou != 0, 'all costs cant be 0' |
| 172 | self.cost_keypoints = cost_keypoints |
| 173 | self.cost_kpvis = cost_kpvis |
| 174 | self.cost_oks = cost_oks |
| 175 | self.focal_alpha = focal_alpha |
| 176 | self.num_body_points = num_body_points |
| 177 | if num_body_points == 17: |
| 178 | self.sigmas = np.array([ |
| 179 | .26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, |
| 180 | 1.07, .87, .87, .89, .89 |
| 181 | ], |
| 182 | dtype=np.float32) / 10.0 |
| 183 | |
| 184 | elif num_body_points == 14: |
| 185 | self.sigmas = np.array([ |
| 186 | .79, .79, .72, .72, .62, .62, 1.07, 1.07, .87, .87, .89, .89, |
| 187 | .79, .79 |
| 188 | ]) / 10.0 |
| 189 | else: |
| 190 | raise ValueError(f'Unsupported keypoints number {num_body_points}') |
| 191 | |
| 192 | @torch.no_grad() |
| 193 | def forward(self, outputs, targets): |
| 194 | bs, num_queries = outputs['pred_logits'].shape[:2] |
| 195 | out_prob = outputs['pred_logits'].flatten(0, 1).sigmoid() |
| 196 | out_bbox = outputs['pred_boxes'].flatten(0, 1) |
| 197 | |
| 198 | |
| 199 | # Also concat the target labels and boxes |
| 200 | tgt_ids = torch.cat([v['labels'] for v in targets]) |
| 201 | tgt_bbox = torch.cat([v['boxes'] for v in targets]) |
| 202 | |
| 203 | # Compute the classification cost. |
| 204 | alpha = self.focal_alpha |
| 205 | gamma = 2.0 |
| 206 | neg_cost_class = (1 - alpha) * (out_prob** |
| 207 | gamma) * (-(1 - out_prob + 1e-8).log()) |
| 208 | pos_cost_class = alpha * ( |
| 209 | (1 - out_prob)**gamma) * (-(out_prob + 1e-8).log()) |
| 210 | cost_class = pos_cost_class[:, tgt_ids] - neg_cost_class[:, tgt_ids] |
| 211 | |
| 212 | # Compute the L1 cost between boxes |
| 213 | cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) |
| 214 | |