diff options
Diffstat (limited to 'Specific/grabber.py')
| -rw-r--r-- | Specific/grabber.py | 113 |
1 files changed, 113 insertions, 0 deletions
diff --git a/Specific/grabber.py b/Specific/grabber.py new file mode 100644 index 0000000..badec7c --- /dev/null +++ b/Specific/grabber.py @@ -0,0 +1,113 @@ +from ctypes import Structure, c_int, POINTER, WINFUNCTYPE, windll, WinError, sizeof +from ctypes.wintypes import BOOL, HWND, RECT, HDC, HBITMAP, HGDIOBJ, DWORD, LONG, WORD, UINT, LPVOID +import numpy as np + +SRCCOPY = 0x00CC0020 +DIB_RGB_COLORS = 0 +BI_RGB = 0 + + +class BITMAPINFOHEADER(Structure): + _fields_ = [('biSize', DWORD), + ('biWidth', LONG), + ('biHeight', LONG), + ('biPlanes', WORD), + ('biBitCount', WORD), + ('biCompression', DWORD), + ('biSizeImage', DWORD), + ('biXPelsPerMeter', LONG), + ('biYPelsPerMeter', LONG), + ('biClrUsed', DWORD), + ('biClrImportant', DWORD)] + + +def err_on_zero_or_null_check(result, func, args): + if not result: + raise WinError() + return args + + +def quick_win_define(name, output, *args, **kwargs): + dllname, fname = name.split('.') + params = kwargs.get('params', None) + if params: + params = tuple([(x, ) for x in params]) + func = (WINFUNCTYPE(output, *args))((fname, getattr(windll, dllname)), params) + err = kwargs.get('err', err_on_zero_or_null_check) + if err: + func.errcheck = err + return func + + +GetClientRect = quick_win_define('user32.GetClientRect', BOOL, HWND, POINTER(RECT), params=(1, 2)) +GetDC = quick_win_define('user32.GetDC', HDC, HWND) +CreateCompatibleDC = quick_win_define('gdi32.CreateCompatibleDC', HDC, HDC) +CreateCompatibleBitmap = quick_win_define('gdi32.CreateCompatibleBitmap', HBITMAP, HDC, c_int, c_int) +ReleaseDC = quick_win_define('user32.ReleaseDC', c_int, HWND, HDC) +DeleteDC = quick_win_define('gdi32.DeleteDC', BOOL, HDC) +DeleteObject = quick_win_define('gdi32.DeleteObject', BOOL, HGDIOBJ) +SelectObject = quick_win_define('gdi32.SelectObject', HGDIOBJ, HDC, HGDIOBJ) +BitBlt = quick_win_define('gdi32.BitBlt', BOOL, HDC, c_int, c_int, c_int, c_int, HDC, c_int, c_int, DWORD) +GetDIBits = quick_win_define('gdi32.GetDIBits', c_int, HDC, HBITMAP, UINT, UINT, LPVOID, POINTER(BITMAPINFOHEADER), UINT) +GetDesktopWindow = quick_win_define('user32.GetDesktopWindow', HWND) + + +class Grabber(object): + def __init__(self, window=None, with_alpha=False, bbox=None): + window = window or GetDesktopWindow() + self.window = window + rect = GetClientRect(window) + self.width = rect.right - rect.left + self.height = rect.bottom - rect.top + if bbox: + bbox = [bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]] + if not bbox[2] or not bbox[3]: + bbox[2] = self.width - bbox[0] + bbox[3] = self.height - bbox[1] + self.x, self.y, self.width, self.height = bbox + else: + self.x = 0 + self.y = 0 + self.windowDC = GetDC(window) + self.memoryDC = CreateCompatibleDC(self.windowDC) + self.bitmap = CreateCompatibleBitmap(self.windowDC, self.width, self.height) + self.bitmapInfo = BITMAPINFOHEADER() + self.bitmapInfo.biSize = sizeof(BITMAPINFOHEADER) + self.bitmapInfo.biPlanes = 1 + self.bitmapInfo.biBitCount = 32 if with_alpha else 24 + self.bitmapInfo.biWidth = self.width + self.bitmapInfo.biHeight = -self.height + self.bitmapInfo.biCompression = BI_RGB + self.bitmapInfo.biSizeImage = 0 + self.channels = 4 if with_alpha else 3 + self.closed = False + + + def __del__(self): + try: + self.close() + except: + pass + + + def close(self): + if self.closed: + return + ReleaseDC(self.window, self.windowDC) + DeleteDC(self.memoryDC) + DeleteObject(self.bitmap) + self.closed = True + + + def grab(self, output=None): + if self.closed: + raise ValueError('Grabber already closed') + if output is None: + output = np.empty((self.height, self.width, self.channels), dtype='uint8') + else: + if output.shape != (self.height, self.width, self.channels): + raise ValueError('Invalid output dimentions') + SelectObject(self.memoryDC, self.bitmap) + BitBlt(self.memoryDC, 0, 0, self.width, self.height, self.windowDC, self.x, self.y, SRCCOPY) + GetDIBits(self.memoryDC, self.bitmap, 0, self.height, output.ctypes.data, self.bitmapInfo, DIB_RGB_COLORS) + return output
\ No newline at end of file |