PyTorch with yolov5：颜色通道和结果显示答案

【问题标题】：PyTorch with yolov5: color channel and result displayPyTorch with yolov5：颜色通道和结果显示
【发布时间】：2021-03-29 11:49:49
【问题描述】：

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

以下作品：

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() 那样在单个窗口中显示带有边界框的结果模型图像吗？（result.show() 为每一帧打开一个新的 Paint 流程实例）？如何将此结果图像保存到磁盘并找到更多关于如何与model 对象交互的文档？

【问题讨论】：

标签： python opencv pytorch yolo

【解决方案1】：

我相信 cvtColor 操作应该与 YOLOv5 PyTorch Hub tutorial 中提供的通道顺序反转相同。这将在两个测试环境（colab notebook python 3.6 和 MacOS python 3.9）中返回 True

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))

【讨论】：

【解决方案2】：

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

【讨论】：