Test

non-max-suppression-in-pytorch-3.py

from torchvision.utils import draw_bounding_boxes
from torchvision.transforms.functional import to_tensor
from typing import List

def plot_bboxes(img : Image.Image, bboxes: torch.Tensor, *args, **kwargs) -> plt.Figure:
    w, h = img.size
    # from [0, 1] to image size
    bboxes = bboxes.clone()
    bboxes[...,0] *= h
    bboxes[...,1] *= w
    bboxes[...,2] *= h
    bboxes[...,3] *= w
    fig = plt.figure()
    img_with_bboxes = draw_bounding_boxes((to_tensor(img) * 255).to(torch.uint8), bboxes, *args, **kwargs, width=4)
    return plt.imshow(img_with_bboxes.permute(1,2,0).numpy())

plot_bboxes(img, original_bboxes, labels=["head", "mic"])