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)
| 558 | |
| 559 | @torch.no_grad() |
| 560 | def forward(self, outputs, target_sizes): |
| 561 | """ Perform the computation |
| 562 | Parameters: |
| 563 | outputs: raw outputs of the model |
| 564 | target_sizes: tensor of dimension [batch_size x 2] containing the size of each images of the batch |
| 565 | For evaluation, this must be the original image size (before any data augmentation) |
| 566 | For visualization, this should be the image size after data augment, but before padding |
| 567 | """ |
| 568 | out_logits, out_bbox = outputs['pred_logits'], outputs['pred_boxes'] |
| 569 | |
| 570 | assert len(out_logits) == len(target_sizes) |
| 571 | assert target_sizes.shape[1] == 2 |
| 572 | |
| 573 | prob = out_logits.sigmoid() |
| 574 | |
| 575 | # topk_values, topk_indexes = torch.topk(prob.view(out_logits.shape[0], -1), 100, dim=1) |
| 576 | # scores = topk_values |
| 577 | # topk_boxes = topk_indexes // out_logits.shape[2] |
| 578 | # labels = topk_indexes % out_logits.shape[2] |
| 579 | # boxes = box_ops.box_cxcywh_to_xyxy(out_bbox) |
| 580 | # boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1,1,4)) |
| 581 | |
| 582 | scores, labels = prob[..., 1:2].max(-1) |
| 583 | labels = labels + 1 |
| 584 | boxes = box_ops.box_cxcywh_to_xyxy(out_bbox) |
| 585 | |
| 586 | # and from relative [0, 1] to absolute [0, height] coordinates |
| 587 | img_h, img_w = target_sizes.unbind(1) |
| 588 | scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1) |
| 589 | boxes = boxes * scale_fct[:, None, :] |
| 590 | |
| 591 | results = [{'scores': s, 'labels': l, 'boxes': b} for s, l, b in zip(scores, labels, boxes)] |
| 592 | |
| 593 | return results |
| 594 | |
| 595 | |
| 596 | class MLP(nn.Module): |