时间戳

首页 > 解决方案 > PyTorch with yolov5：颜色通道和结果显示

python - PyTorch with yolov5：颜色通道和结果显示

问题描述

我有一个脚本可以抓取应用程序的屏幕截图并显示它。它在我的机器上运行得非常好，就像一个大约 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?

我的问题：

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

标签： pythonopencvpytorchyolo

解决方案

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

推荐阅读