如何得到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 上有关此主题的现有信息

查看本节中yolov5 的官方 Pytorch Hub Wiki，Base64 Results我们可以找到有关如何使用的信息render，以及 yolov5 和 REST API 的一些方便信息以及实现此功能的原因。

查看代码以了解此行为

type(results)如果看一下我们得到的结果类型，我们可以在yolov5 的 GitHub 页面<class 'models.common.Detections'>找到它。

该类具有 this （render、show等）方法，这些方法_run使用不同的参数运行该方法。该render方法将导致该方法_run用检测到的图像覆盖所有图像：

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