This module converts the model's output into the format expected by the coco api
| 326 | |
| 327 | |
| 328 | class TrackerPostProcess(nn.Module): |
| 329 | """ This module converts the model's output into the format expected by the coco api""" |
| 330 | def __init__(self): |
| 331 | super().__init__() |
| 332 | |
| 333 | @torch.no_grad() |
| 334 | def forward(self, track_instances: Instances, target_size) -> Instances: |
| 335 | """ Perform the computation |
| 336 | Parameters: |
| 337 | outputs: raw outputs of the model |
| 338 | target_sizes: tensor of dimension [batch_size x 2] containing the size of each images of the batch |
| 339 | For evaluation, this must be the original image size (before any data augmentation) |
| 340 | For visualization, this should be the image size after data augment, but before padding |
| 341 | """ |
| 342 | out_logits = track_instances.pred_logits |
| 343 | out_bbox = track_instances.pred_boxes |
| 344 | |
| 345 | prob = out_logits.sigmoid() |
| 346 | # prob = out_logits[...,:1].sigmoid() |
| 347 | scores, labels = prob.max(-1) |
| 348 | |
| 349 | # convert to [x0, y0, x1, y1] format |
| 350 | boxes = box_ops.box_cxcywh_to_xyxy(out_bbox) |
| 351 | # and from relative [0, 1] to absolute [0, height] coordinates |
| 352 | img_h, img_w = target_size |
| 353 | scale_fct = torch.Tensor([img_w, img_h, img_w, img_h]).to(boxes) |
| 354 | boxes = boxes * scale_fct[None, :] |
| 355 | |
| 356 | track_instances.boxes = boxes |
| 357 | track_instances.scores = scores |
| 358 | track_instances.labels = labels |
| 359 | # track_instances.remove('pred_logits') |
| 360 | # track_instances.remove('pred_boxes') |
| 361 | return track_instances |
| 362 | |
| 363 | |
| 364 | def _get_clones(module, N): |