'python screengrab problem (on yolov5, pytorch)

im trying to get a screenshot of a window as fast and then inference on yolov5 it works but sometimes it doesnt detect very well compared to using detect.py on the same image. i think its probably because of img shape or array but i dont know where or how to edit those to make it work. can anyone help me with this please?

import torch
import numpy as np
import win32gui
import win32ui
import win32con



w = 800 # set this
h = 600 # set this
bmpfilenamename = "color.bmp" #set this
windowname = 'put windowname'


def screenshot():
    hwnd = win32gui.FindWindow(None, windowname)
    wDC = win32gui.GetWindowDC(hwnd)
    dcObj=win32ui.CreateDCFromHandle(wDC)
    cDC=dcObj.CreateCompatibleDC()
    dataBitMap = win32ui.CreateBitmap()
    dataBitMap.CreateCompatibleBitmap(dcObj, w, h)
    cDC.SelectObject(dataBitMap)
    cDC.BitBlt((0,0),(w, h) , dcObj, (0,0), win32con.SRCCOPY)

#save the screenshot
#dataBitMap.SaveBitmapFile(cDC, bmpfilenamename)

signedIntsArray = dataBitMap.GetBitmapBits(True)
img = np.frombuffer(signedIntsArray, dtype='uint8')
img.shape = (h,w,4)

# Free Resources
dcObj.DeleteDC()
cDC.DeleteDC()
win32gui.ReleaseDC(hwnd, wDC)
win32gui.DeleteObject(dataBitMap.GetHandle())

#img = img[..., ::-1]
#img = np.ascontiguousarray(img)

return img

#load
model = torch.hub.load('./', 'custom', path='yolov5s.pt', source='local')


#inference
test = screenshot()
results = model(test)
boxes = results.pandas().xyxy[0]
print (boxes)

edit : i figured you can do it by changing the code from this

results = model(test) 

to this

results = model(cv.cvtColor(test, cv.COLOR_BGR2RGB))

but isnt this code supposed to do the same? for some reason this one wont work

img[:    ,:    ,::-1]


Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source