如何从 pytoch 中 yolov5 给出的结果中检索结果图像作为矩阵（numpy 数组）？

Question

我一直在学习如何使用pytorch实现预训练的yolo，我想使用openCV的cv2.imshow()方法显示输出图像。 输出图像可以使用 .show() 函数显示并使用 .save() 函数保存，但是我想使用 cv2.imshow() 显示它，为此我需要 numpy 数组形式的图像。 我不知道我们如何做到这一点，也不知道这是否可能。 这是它的代码。

import torch

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
imgs = ['img.png']  # batch of images

results = model(imgs)

results.print()
results.show()  # or .save(), shows/saves the same image with bounding boxes around detected objects

# Show 'results' using openCV's cv2.imshow() method?

results.xyxy[0]  # img1 predictions (tensor)
print(results.pandas().xyxy[0])  # img1 predictions (pandas)

解决这个问题的一个更长的方法是在图像中检测到的对象上自己创建边界框并显示它，但请考虑我的懒惰：p。

Answer 1

我和你一样懒:)，你可以显示边界框，而不需要手动绘制它们。当您调用

results.save()

时，它会将带有框的图像版本保存到此文件夹“runs/detect/exp/”，然后您可以使用 cv2 显示该图像。

results.save()
import cv2
img = cv2.imread("runs/detect/exp/zidane.jpg")

#cv2.imgshow does not work on Google collab so this is a work around. 
# You should get the same results if you use cv2.imgshow 
from google.colab.patches import cv2_imshow 
cv2_imshow(img)

Answer 2

有同样的问题，所以我写了一个小方法来快速绘制图像而不保存它。

  def drawRectangles(image, dfResults):
    for index, row in dfResults.iterrows():
      print( (row['xmin'], row['ymin']))
      image = cv2.rectangle(image, (row['xmin'], row['ymin']), (row['xmax'], row['ymax']), (255, 0, 0), 2)
    cv2_imshow(image)

_

    results = model(image)
    dfResults = results.pandas().xyxy[0]
    self.drawRectangles(image, dfResults[['xmin', 'ymin', 'xmax','ymax']].astype(int))

Answer 3

只需通过 cv2 打开图像，然后通过结果数组中的给定点绘制矩形来广告矩形？

import cv2

img = cv2.imread("image_path")

链接如何使用 cv2 绘制矩形

@ChristophRackwitz

你有它

cv2.rectangle(image, start_point, end_point, color, thickness)

Answer 4

当然有一种简单的方法可以将预测结果获取为 numpy 数组。假设测试图像的路径是'path/to_image'，那么。

`结果=模型（'路径/to_image'）

predicted_numpy_image=结果[0].plot()

cv2.imshow('winname',predicted_numpy_image) `