我正在尝试加载 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)
谢谢。
最佳答案
因为@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)
关于opencv - 如何得到YOLOv5模型的预测图像?,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/71995836/