如何得到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

一个快速的解决方法是使用“pred”中的“imgs”对象，如下所示：

predict_image = model(image)
im_rgb = cv2.cvtColor(predict_image.imgs[0], cv2.COLOR_BGR2RGB) # Because of OpenCV reading images as BGR
cv2_imshow(im_rgb)

希望这对你有所帮助，祝你有美好的一天。

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

How to get a predicted image of YOLOv5 model?

opencv

pytorch

yolo