Files
seshat-tts/src/seshat_tts/capture.py
T
cbartos 75fc1afa53
CI / Tests (3.10) (push) Has been cancelled
CI / Tests (3.13) (push) Has been cancelled
seshat-tts
2026-05-22 05:54:01 -04:00

117 lines
3.5 KiB
Python

from __future__ import annotations
import ctypes
from dataclasses import dataclass
import mss
from PIL import Image
import win32gui
import win32ui
from .config import Rect
@dataclass(frozen=True, slots=True)
class MonitorInfo:
index: int
left: int
top: int
width: int
height: int
@property
def label(self) -> str:
return f"{self.index}: {self.width}x{self.height} at {self.left},{self.top}"
def list_monitors() -> list[MonitorInfo]:
with mss.mss() as sct:
return [
MonitorInfo(
index=index,
left=int(monitor["left"]),
top=int(monitor["top"]),
width=int(monitor["width"]),
height=int(monitor["height"]),
)
for index, monitor in enumerate(sct.monitors)
if index != 0
]
def capture_absolute_region(left: int, top: int, width: int, height: int) -> Image.Image:
with mss.mss() as sct:
grab = {
"left": left,
"top": top,
"width": width,
"height": height,
}
shot = sct.grab(grab)
return Image.frombytes("RGB", shot.size, shot.rgb)
def capture_monitor_region(monitor_index: int, rect: Rect) -> Image.Image:
with mss.mss() as sct:
if monitor_index <= 0 or monitor_index >= len(sct.monitors):
raise ValueError(f"Monitor {monitor_index} is not available.")
monitor = sct.monitors[monitor_index]
return capture_absolute_region(
int(monitor["left"]) + rect.left,
int(monitor["top"]) + rect.top,
rect.width,
rect.height,
)
def capture_window_region(hwnd: int, rect: Rect) -> Image.Image:
image = capture_window(hwnd)
if rect.left < 0 or rect.top < 0 or rect.width <= 0 or rect.height <= 0:
raise ValueError("Capture region must be inside the selected window.")
if rect.left + rect.width > image.width or rect.top + rect.height > image.height:
raise ValueError("Capture region is outside the selected window. Select the region again in window mode.")
return image.crop((rect.left, rect.top, rect.left + rect.width, rect.top + rect.height))
def capture_window(hwnd: int) -> Image.Image:
left, top, right, bottom = win32gui.GetWindowRect(hwnd)
width = right - left
height = bottom - top
if width <= 0 or height <= 0:
raise ValueError("Selected window has no capturable size.")
hwnd_dc = win32gui.GetWindowDC(hwnd)
source_dc = win32ui.CreateDCFromHandle(hwnd_dc)
memory_dc = source_dc.CreateCompatibleDC()
bitmap = win32ui.CreateBitmap()
bitmap.CreateCompatibleBitmap(source_dc, width, height)
memory_dc.SelectObject(bitmap)
try:
result = _print_window(hwnd, memory_dc.GetSafeHdc(), 2)
if result != 1:
result = _print_window(hwnd, memory_dc.GetSafeHdc(), 0)
if result != 1:
raise RuntimeError("PrintWindow failed for the selected window.")
info = bitmap.GetInfo()
bits = bitmap.GetBitmapBits(True)
return Image.frombuffer(
"RGB",
(info["bmWidth"], info["bmHeight"]),
bits,
"raw",
"BGRX",
0,
1,
).copy()
finally:
win32gui.DeleteObject(bitmap.GetHandle())
memory_dc.DeleteDC()
source_dc.DeleteDC()
win32gui.ReleaseDC(hwnd, hwnd_dc)
def _print_window(hwnd: int, hdc: int, flags: int) -> int:
return int(ctypes.windll.user32.PrintWindow(hwnd, hdc, flags))