PyTorch with yolov5：颜色通道和结果展示

Question

我有一个脚本可以抓取应用程序的屏幕截图并显示它。它在我的机器上运行得非常好，就像一个大约 60FPS 的视频。现在我想在这些帧上使用 yolov5 对象检测模型，使用 TorchHub，按照建议。

以下作品：

import os
os.getcwd()
from PIL import ImageGrab
import numpy as np
import cv2
import pyautogui
import win32gui
import time
from mss import mss
from PIL import Image
import tempfile
os.system('calc')
sct = mss()
xx=1
tstart = time.time()
while xx<10000:
    hwnd = win32gui.FindWindow(None, 'Calculator')
    left_x, top_y, right_x, bottom_y = win32gui.GetWindowRect(hwnd)
    #screen = np.array(ImageGrab.grab( bbox = (left_x, top_y, right_x, bottom_y ) ) )
    bbox = {'top': top_y, 'left': left_x, 'width': right_x-left_x, 'height':bottom_y-top_y }
    screen = sct.grab(bbox)
    scr = np.array(screen)
    
    cv2.imshow('window', scr)
    if cv2.waitKey(25) & 0xFF == ord('q'):
        cv2.destroyAllWindows()
        break
    xx+=1
cv2.destroyAllWindows()
tend = time.time()
print(xx/(tend-tstart))
print((tend-tstart))
os.system('taskkill /f /im calculator.exe')

下面我尝试import torch并使用我之前训练的模型，

screen = sct.grab(bbox)
scr = np.array(screen)    
result = model(scr, size=400)  
result.save("test.png") #this gives a TypeError: save() takes 1 positional argument but 2 were given
result.show() #this opens a new Paint instance for every frame instead of keeping the same window. 
# The shown image is also in a wrong color channel
scr = cv2.imread("test.png")
# How can I use the `result` as argument to cv2.imshow(),
# without saving to disk if possible?

我的问题：

cv2.imshow()

result.show() 显示的图像颜色通道错误，我如何确保提供给 model 的图像在正确的通道上？
与训练验证相比，分类和检测的性能急剧下降，可能是因为 1?
您知道我如何像 cv2.imshow() 那样在单个 window 中显示带有边界框的结果模型图像吗？（result.show() 为每一帧打开一个新的 Paint 进程实例） ?如何将此结果图像保存到磁盘并找到更多关于如何与 model 对象交互的一般文档？

Answer 1

以下有效：result = model(cv2.cvtColor(scr, cv2.COLOR_BGR2RGB), size=400) 这解决了准确性问题并且 model.save() 具有当前不可更改的预定义输出名称，它不带任何参数。 model.show() 将正确的颜色通道作为输入时显示正确的颜色通道输出。

Answer 2

我相信 cvtColor 操作应该与 YOLOv5 PyTorch Hub tutorial 中显示的提供的通道顺序反转相同。这个returnsTrue在两个环境下测试（colab notebookpython3.6和MacOSpython3.9）

import cv2
import numpy as np

file = 'data/images/bus.jpg'
im1 = cv2.imread(file)[:, :, ::-1]
im2 = cv2.cvtColor(cv2.imread(file), cv2.COLOR_BGR2RGB)
print(np.allclose(im1, im2))

PyTorch with yolov5：颜色通道和结果展示

PyTorch with yolov5: color channel and result display

python

opencv

pytorch

yolo