opencv - 如何得到YOLOv5模型的预测图像？

Question

我正在尝试加载 YOLOv5 模型并使用它来预测特定图像。我的问题是我想在我的应用程序中显示带有边界框的预测图像，因此我需要直接从 PyTorch 的预测方法中获取它以在我的应用程序中显示。

    model = torch.hub.load('yolov5', 'custom', path=model_name, force_reload=True, 
    source='local')
    pred = model(image)
    pred.show() #show image but can't assign to a variable
    pred.save() #save image to runs\detect\exp

我想要这样的东西:

    predict_image = model(image)
    cv2.imshow('Predict', predict_image)

谢谢。

Answer 1

因为@mak13 的回答对我不起作用，我想了解这种行为，这是我的解决方案:

简单获取图像作为变量

import torch
import pathlib

img_path = pathlib.Path("test_img.jpg")

model = torch.hub.load('ultralytics/yolov5', 'yolov5n')
results = model(img_path)
r_img = results.render() # returns a list with the images as np.array
img_with_boxes = r_img[0] # image with boxes as np.array

GitHub 上有关此主题的现有信息

看官方Pytorch Hub Wiki from yolov5在 Base64 Results 部分我们找到有关如何使用 render 的信息以及有关 yolov5 和 REST API 的一些实用信息，以及实现这一点的原因。

查看代码以了解此行为

如果查看结果类型 type(results)我们得到 <class 'models.common.Detections'>我们可以在这里找到 yolov5's GitHub page .

此类具有运行 render 的此( show 、 _run 等)方法。具有不同参数的方法。 render方法将导致 _run用检测到的图像覆盖所有图像的方法:

if render:
    self.ims[i] = np.asarray(im)