在pytorch中的mobilenet_v2,如何获得盒子和标签?

我想看看由 mobilenet_v2 模型,这是我迄今为止的代码:

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)

我试着使用这个代码来查看带注释的图像

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

但它返回: annotated_frame=输出[0].plot() AttributeError:“Tensor”对象没有属性“plot”。你的意思是:“漂浮”吗?

事实上,经过检查 输出 似乎是浮子的张量。。。

那么,关于如何获得带注释的框架,有什么想法吗?

谢谢!