mobilenet_v2 in pytorch, how to get boxes and labels?

Question

I want to see the boxes and labels of the predictions made by the mobilenet_v2 model, here is my code so far:

import time

import torch
import numpy as np
from torchvision import models, transforms

import cv2
from PIL import Image


cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 224)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 224)
cap.set(cv2.CAP_PROP_FPS, 36)

preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

net = models.quantization.mobilenet_v2(pretrained=True, quantize=False)
net = torch.jit.script(net)

started = time.time()
last_logged = time.time()
frame_count = 0

with torch.no_grad():
    while True:

        # read frame
        ret, image = cap.read()
        if not ret:
            raise RuntimeError("failed to read frame")

        # convert opencv output from BGR to RGB
        image = image[:, :, [2, 1, 0]]
        permuted = image

        # preprocess
        input_tensor = preprocess(image)

        # create a mini-batch as expected by the model
        input_batch = input_tensor.unsqueeze(0)

        # run model
        output = net(input_batch)

        # I tried the code below in order to retrieve annotated images
       
        annotated_frame = output[0].plot()
        cv2.imshow("YOLOv8n Inference", annotated_frame)

I have tried to use this code in order to see the annotated images

annotated_frame = output[0].plot()
cv2.imshow("YOLOv8n Inference", annotated_frame)

but it returns: annotated_frame = output[0].plot() AttributeError: 'Tensor' object has no attribute 'plot'. Did you mean: 'float'?

indeed, after inspection, the output seems be a tensor of floats...

So, any idea on how to obtain the annotated frames ?

Thanks a lot!