window_capture.py

import numpy as np
import win32gui, win32ui, win32con
from threading import Thread, Lock

class WindowCapture:

    # Threading properties
    stopped = True
    lock = None
    screenshot = None

    # properties
    w, h = 0, 0
    hwnd = None
    cropped_x = 0
    cropped_y = 0
    offset_x = 0
    offset_y = 0

    # constructor
    def __init__(self, window_name=None):
        # Create thread lock
        self.lock = Lock()

        # Find window to capture
        if window_name is None:
            self.hwnd = win32gui.GetDesktopWindow()
        else:
            self.hwnd = win32gui.FindWindow(None, window_name)
            if not self.hwnd:
                raise Exception('Window not found: {}'.format(window_name))

        # Get window size
        window_rect = win32gui.GetWindowRect(self.hwnd)
        self.w = window_rect[2] - window_rect[0]
        self.h = window_rect[3] - window_rect[1]

        # Remove border and title bar
        border_pixels = 5
        titlebar_pixels = 27
        self.w = self.w - (border_pixels * 2)
        self.h = self.h - titlebar_pixels - border_pixels
        self.cropped_x = border_pixels
        self.cropped_y = titlebar_pixels

        # Set cropped coords so we can translate screenshot images to actual positions
        self.offset_x = window_rect[0] + self.cropped_x
        self.offset_y = window_rect[1] + self.cropped_y

    def get_screenshot(self):

        # get the window image data
        wDC = win32gui.GetWindowDC(self.hwnd)
        dcObj = win32ui.CreateDCFromHandle(wDC)
        cDC = dcObj.CreateCompatibleDC()
        dataBitMap = win32ui.CreateBitmap()
        dataBitMap.CreateCompatibleBitmap(dcObj, self.w, self.h)
        cDC.SelectObject(dataBitMap)
        cDC.BitBlt((0, 0), (self.w, self.h), dcObj, (self.cropped_x, self.cropped_y), win32con.SRCCOPY)

        # save the image as a bitmap file
        #dataBitMap.SaveBitmapFile(cDC, 'yolov7/inference/images/osrs.bmp')

        # convert the raw data into a format opencv can read
        signedIntsArray = dataBitMap.GetBitmapBits(True)
        img = np.fromstring(signedIntsArray, dtype='uint8')
        img.shape = (self.h, self.w, 4)

        # free resources
        dcObj.DeleteDC()
        cDC.DeleteDC()
        win32gui.ReleaseDC(self.hwnd, wDC)
        win32gui.DeleteObject(dataBitMap.GetHandle())

        # Drop alpha channel data (transparency) of open cv match template
        img = img[...,:3]

        # Fix errors from opencv_points.py
        img = np.ascontiguousarray(img)


        return img

    # Translate a pixel position on a screenshot image to a pixel position on the screen.
    # pos = (x, y)
    # WARNING: if you move the window being captured after execution is started, this will
    # return incorrect coordinates, because the window position is only calculated in the
    # __init__ constructor
    def get_screen_pos(self, pos):
        return (pos[0] + self.offset_x, pos[1] + self.offset_y)

    # Threading methods

    def start(self):
        self.stopped = False
        t = Thread(target=self.run)
        t.start()

    def stop(self):
        self.stopped = True

    def run(self):

        while not self.stopped:
            # Get screenshot
            screenshot = self.get_screenshot()
            # Lock thread while getting result
            self.lock.acquire()
            self.screenshot = screenshot
            self.lock.release()