This module converts the model's output into the format expected by the coco api
| 243 | |
| 244 | |
| 245 | class PostProcess(nn.Module): |
| 246 | """ This module converts the model's output into the format expected by the coco api""" |
| 247 | @torch.no_grad() |
| 248 | def forward(self, outputs, target_sizes): |
| 249 | """ Perform the computation |
| 250 | Parameters: |
| 251 | outputs: raw outputs of the model |
| 252 | target_sizes: tensor of dimension [batch_size x 2] containing the size of each images of the batch |
| 253 | For evaluation, this must be the original image size (before any data augmentation) |
| 254 | For visualization, this should be the image size after data augment, but before padding |
| 255 | """ |
| 256 | out_logits, out_bbox = outputs['pred_logits'], outputs['pred_boxes'] |
| 257 | |
| 258 | assert len(out_logits) == len(target_sizes) |
| 259 | assert target_sizes.shape[1] == 2 |
| 260 | |
| 261 | prob = F.softmax(out_logits, -1) |
| 262 | scores, labels = prob[..., :-1].max(-1) |
| 263 | |
| 264 | # convert to [x0, y0, x1, y1] format |
| 265 | boxes = box_ops.box_cxcywh_to_xyxy(out_bbox) |
| 266 | # and from relative [0, 1] to absolute [0, height] coordinates |
| 267 | img_h, img_w = target_sizes.unbind(1) |
| 268 | scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1) |
| 269 | boxes = boxes * scale_fct[:, None, :] |
| 270 | |
| 271 | results = [{'scores': s, 'labels': l, 'boxes': b} for s, l, b in zip(scores, labels, boxes)] |
| 272 | |
| 273 | return results |
| 274 | |
| 275 | |
| 276 |