| from __future__ import annotations |
|
|
| import torch |
| from torch import Tensor |
|
|
| from .boxes import box_cxcywh_to_xyxy |
|
|
|
|
| @torch.no_grad() |
| def decode_predictions( |
| outputs: dict[str, Tensor], |
| image_sizes: list[tuple[int, int]], |
| confidence: float = 0.25, |
| top_k: int = 300, |
| ) -> list[dict[str, Tensor]]: |
| logits = outputs["pred_logits"].sigmoid() |
| boxes = box_cxcywh_to_xyxy(outputs["pred_boxes"]).clamp(0.0, 1.0) |
| results = [] |
| for index, (height, width) in enumerate(image_sizes): |
| scores, labels = logits[index].max(dim=-1) |
| keep = scores >= confidence |
| if keep.sum() > top_k: |
| selected = scores.masked_fill(~keep, -1).topk(top_k).indices |
| else: |
| selected = torch.where(keep)[0] |
| selected_boxes = boxes[index, selected].clone() |
| selected_boxes[:, [0, 2]] *= width |
| selected_boxes[:, [1, 3]] *= height |
| results.append( |
| {"scores": scores[selected], "labels": labels[selected], "boxes": selected_boxes} |
| ) |
| return results |
|
|
