Perform the computation Parameters: outputs: raw outputs of the model target_sizes: tensor of dimension [batch_size x 2] containing the size of each images of the batch For evaluation, this must be the original image size (before any data au
(self, outputs, target_sizes)
| 447 | |
| 448 | @torch.no_grad() |
| 449 | def forward(self, outputs, target_sizes): |
| 450 | """ Perform the computation |
| 451 | Parameters: |
| 452 | outputs: raw outputs of the model |
| 453 | target_sizes: tensor of dimension [batch_size x 2] containing the size of each images of the batch |
| 454 | For evaluation, this must be the original image size (before any data augmentation) |
| 455 | For visualization, this should be the image size after data augment, but before padding |
| 456 | """ |
| 457 | out_logits, out_bbox = outputs['pred_logits'], outputs['pred_boxes'] |
| 458 | track_logits, track_bbox = outputs['tracking_logits'], outputs['tracking_boxes'] |
| 459 | |
| 460 | assert len(out_logits) == len(target_sizes) |
| 461 | assert target_sizes.shape[1] == 2 |
| 462 | |
| 463 | prob = out_logits.sigmoid() |
| 464 | track_prob = track_logits.sigmoid() |
| 465 | # topk_values, topk_indexes = torch.topk(prob.view(out_logits.shape[0], -1), 100, dim=1) |
| 466 | # scores = topk_values |
| 467 | # topk_boxes = topk_indexes // out_logits.shape[2] |
| 468 | # labels = topk_indexes % out_logits.shape[2] |
| 469 | # boxes = box_ops.box_cxcywh_to_xyxy(out_bbox) |
| 470 | # boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1,1,4)) |
| 471 | |
| 472 | scores, labels = prob[..., 1:2].max(-1) |
| 473 | labels = labels + 1 |
| 474 | boxes = box_ops.box_cxcywh_to_xyxy(out_bbox) |
| 475 | |
| 476 | track_scores, track_labels = track_prob[..., 1:2].max(-1) |
| 477 | track_labels = track_labels + 1 |
| 478 | track_boxes = box_ops.box_cxcywh_to_xyxy(track_bbox) |
| 479 | |
| 480 | # and from relative [0, 1] to absolute [0, height] coordinates |
| 481 | img_h, img_w = target_sizes.unbind(1) |
| 482 | scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1) |
| 483 | boxes = boxes * scale_fct[:, None, :] |
| 484 | track_boxes = track_boxes * scale_fct[:, None, :] |
| 485 | |
| 486 | results = [{'scores': s, 'labels': l, 'boxes': b, 'track_scores': ts, 'track_labels': tl, 'track_boxes': tb} |
| 487 | for s, l, b, ts, tl, tb in zip(scores, labels, boxes, track_scores, track_labels, track_boxes)] |
| 488 | |
| 489 | return results |
| 490 | |
| 491 | |
| 492 | class MLP(nn.Module): |