371 lines
11 KiB
Python
371 lines
11 KiB
Python
|
"""
|
||
|
|
Screen capture utilities for GSPro Remote.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import logging
|
||
|
|
from typing import Optional, Tuple, Dict, Any
|
||
|
|
from io import BytesIO
|
||
|
|
import base64
|
||
|
|
|
||
|
|
try:
|
||
|
|
import mss
|
||
|
|
import mss.tools
|
||
|
|
from PIL import Image
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
except ImportError as e:
|
||
|
|
raise ImportError(f"Required screen capture dependencies not installed: {e}")
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class ScreenCapture:
|
||
|
|
"""Manages screen capture operations."""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
"""Initialize screen capture manager."""
|
||
|
|
self.sct = mss.mss()
|
||
|
|
self._monitor_info = None
|
||
|
|
|
||
|
|
def get_monitors(self) -> list[dict]:
|
||
|
|
"""
|
||
|
|
Get information about all available monitors.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
List of monitor information dictionaries
|
||
|
|
"""
|
||
|
|
monitors = []
|
||
|
|
for i, monitor in enumerate(self.sct.monitors):
|
||
|
|
monitors.append(
|
||
|
|
{
|
||
|
|
"index": i,
|
||
|
|
"left": monitor["left"],
|
||
|
|
"top": monitor["top"],
|
||
|
|
"width": monitor["width"],
|
||
|
|
"height": monitor["height"],
|
||
|
|
"is_primary": i == 0, # Index 0 is combined virtual screen
|
||
|
|
}
|
||
|
|
)
|
||
|
|
return monitors
|
||
|
|
|
||
|
|
def get_primary_monitor(self) -> dict:
|
||
|
|
"""
|
||
|
|
Get the primary monitor information.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Primary monitor information
|
||
|
|
"""
|
||
|
|
# Index 1 is typically the primary monitor in mss
|
||
|
|
return self.sct.monitors[1] if len(self.sct.monitors) > 1 else self.sct.monitors[0]
|
||
|
|
|
||
|
|
def capture_screen(self, monitor_index: int = 1) -> np.ndarray:
|
||
|
|
"""
|
||
|
|
Capture the entire screen.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
monitor_index: Index of the monitor to capture (0 for all, 1 for primary)
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Captured screen as numpy array (BGR format)
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
monitor = self.sct.monitors[monitor_index]
|
||
|
|
screenshot = self.sct.grab(monitor)
|
||
|
|
|
||
|
|
# Convert to numpy array (BGR format for OpenCV compatibility)
|
||
|
|
img = np.array(screenshot)
|
||
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)
|
||
|
|
|
||
|
|
return img
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Failed to capture screen: {e}")
|
||
|
|
raise
|
||
|
|
|
||
|
|
def capture_region(self, x: int, y: int, width: int, height: int, monitor_index: int = 1) -> np.ndarray:
|
||
|
|
"""
|
||
|
|
Capture a specific region of the screen.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
x: X coordinate of the region (relative to monitor)
|
||
|
|
y: Y coordinate of the region (relative to monitor)
|
||
|
|
width: Width of the region
|
||
|
|
height: Height of the region
|
||
|
|
monitor_index: Index of the monitor to capture from
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Captured region as numpy array (BGR format)
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
monitor = self.sct.monitors[monitor_index]
|
||
|
|
|
||
|
|
# Define region to capture
|
||
|
|
region = {
|
||
|
|
"left": monitor["left"] + x,
|
||
|
|
"top": monitor["top"] + y,
|
||
|
|
"width": width,
|
||
|
|
"height": height,
|
||
|
|
}
|
||
|
|
|
||
|
|
screenshot = self.sct.grab(region)
|
||
|
|
|
||
|
|
# Convert to numpy array
|
||
|
|
img = np.array(screenshot)
|
||
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)
|
||
|
|
|
||
|
|
return img
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Failed to capture region: {e}")
|
||
|
|
raise
|
||
|
|
|
||
|
|
def capture_window(self, window_title: str) -> Optional[np.ndarray]:
|
||
|
|
"""
|
||
|
|
Capture a specific window by title.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
window_title: Title of the window to capture
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Captured window as numpy array or None if window not found
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
import win32gui
|
||
|
|
|
||
|
|
def enum_window_callback(hwnd, windows):
|
||
|
|
if win32gui.IsWindowVisible(hwnd) and win32gui.IsWindowEnabled(hwnd):
|
||
|
|
window_text = win32gui.GetWindowText(hwnd)
|
||
|
|
if window_title.lower() in window_text.lower():
|
||
|
|
windows.append(hwnd)
|
||
|
|
return True
|
||
|
|
|
||
|
|
windows = []
|
||
|
|
win32gui.EnumWindows(enum_window_callback, windows)
|
||
|
|
|
||
|
|
if not windows:
|
||
|
|
logger.warning(f"Window not found: {window_title}")
|
||
|
|
return None
|
||
|
|
|
||
|
|
# Get window rectangle
|
||
|
|
hwnd = windows[0]
|
||
|
|
rect = win32gui.GetWindowRect(hwnd)
|
||
|
|
x, y, x2, y2 = rect
|
||
|
|
width = x2 - x
|
||
|
|
height = y2 - y
|
||
|
|
|
||
|
|
# Capture the window region
|
||
|
|
return self.capture_region(x, y, width, height, monitor_index=0)
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Failed to capture window: {e}")
|
||
|
|
return None
|
||
|
|
|
||
|
|
def image_to_base64(self, image: np.ndarray, quality: int = 85, format: str = "JPEG") -> str:
|
||
|
|
"""
|
||
|
|
Convert an image array to base64 string.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
image: Image as numpy array (BGR format)
|
||
|
|
quality: JPEG quality (1-100)
|
||
|
|
format: Image format (JPEG, PNG)
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Base64 encoded image string
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
# Convert BGR to RGB for PIL
|
||
|
|
rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||
|
|
pil_image = Image.fromarray(rgb_image)
|
||
|
|
|
||
|
|
# Save to bytes
|
||
|
|
buffer = BytesIO()
|
||
|
|
if format.upper() == "JPEG":
|
||
|
|
pil_image.save(buffer, format=format, quality=quality, optimize=True)
|
||
|
|
else:
|
||
|
|
pil_image.save(buffer, format=format)
|
||
|
|
|
||
|
|
# Encode to base64
|
||
|
|
buffer.seek(0)
|
||
|
|
base64_string = base64.b64encode(buffer.getvalue()).decode("utf-8")
|
||
|
|
|
||
|
|
return base64_string
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Failed to convert image to base64: {e}")
|
||
|
|
raise
|
||
|
|
|
||
|
|
def resize_image(self, image: np.ndarray, width: Optional[int] = None, height: Optional[int] = None) -> np.ndarray:
|
||
|
|
"""
|
||
|
|
Resize an image while maintaining aspect ratio.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
image: Image as numpy array
|
||
|
|
width: Target width (None to calculate from height)
|
||
|
|
height: Target height (None to calculate from width)
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Resized image as numpy array
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
h, w = image.shape[:2]
|
||
|
|
|
||
|
|
if width and not height:
|
||
|
|
# Calculate height maintaining aspect ratio
|
||
|
|
height = int(h * (width / w))
|
||
|
|
elif height and not width:
|
||
|
|
# Calculate width maintaining aspect ratio
|
||
|
|
width = int(w * (height / h))
|
||
|
|
elif not width and not height:
|
||
|
|
# No resize needed
|
||
|
|
return image
|
||
|
|
|
||
|
|
return cv2.resize(image, (width, height), interpolation=cv2.INTER_AREA)
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Failed to resize image: {e}")
|
||
|
|
raise
|
||
|
|
|
||
|
|
def get_resolution_preset(self, preset: str) -> Tuple[int, int]:
|
||
|
|
"""
|
||
|
|
Get width and height for a resolution preset.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
preset: Resolution preset (e.g., '720p', '1080p', '480p')
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Tuple of (width, height)
|
||
|
|
"""
|
||
|
|
presets = {
|
||
|
|
"480p": (854, 480),
|
||
|
|
"540p": (960, 540),
|
||
|
|
"720p": (1280, 720),
|
||
|
|
"900p": (1600, 900),
|
||
|
|
"1080p": (1920, 1080),
|
||
|
|
"1440p": (2560, 1440),
|
||
|
|
"2160p": (3840, 2160),
|
||
|
|
"4k": (3840, 2160),
|
||
|
|
}
|
||
|
|
|
||
|
|
return presets.get(preset.lower(), (1280, 720))
|
||
|
|
|
||
|
|
def close(self):
|
||
|
|
"""Close the screen capture resources."""
|
||
|
|
if hasattr(self, "sct"):
|
||
|
|
self.sct.close()
|
||
|
|
|
||
|
|
|
||
|
|
# Global screen capture instance
|
||
|
|
_screen_capture: Optional[ScreenCapture] = None
|
||
|
|
|
||
|
|
|
||
|
|
def get_screen_capture() -> ScreenCapture:
|
||
|
|
"""
|
||
|
|
Get the global screen capture instance.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
ScreenCapture instance
|
||
|
|
"""
|
||
|
|
global _screen_capture
|
||
|
|
if _screen_capture is None:
|
||
|
|
_screen_capture = ScreenCapture()
|
||
|
|
return _screen_capture
|
||
|
|
|
||
|
|
|
||
|
|
def capture_screen(monitor_index: int = 1) -> np.ndarray:
|
||
|
|
"""
|
||
|
|
Capture the entire screen.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
monitor_index: Index of the monitor to capture
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Captured screen as numpy array
|
||
|
|
"""
|
||
|
|
return get_screen_capture().capture_screen(monitor_index)
|
||
|
|
|
||
|
|
|
||
|
|
def capture_region(x: int, y: int, width: int, height: int) -> np.ndarray:
|
||
|
|
"""
|
||
|
|
Capture a specific region of the screen.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
x: X coordinate of the region
|
||
|
|
y: Y coordinate of the region
|
||
|
|
width: Width of the region
|
||
|
|
height: Height of the region
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Captured region as numpy array
|
||
|
|
"""
|
||
|
|
return get_screen_capture().capture_region(x, y, width, height)
|
||
|
|
|
||
|
|
|
||
|
|
def get_screen_size() -> Tuple[int, int]:
|
||
|
|
"""
|
||
|
|
Get the primary screen size.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Tuple of (width, height)
|
||
|
|
"""
|
||
|
|
monitor = get_screen_capture().get_primary_monitor()
|
||
|
|
return monitor["width"], monitor["height"]
|
||
|
|
|
||
|
|
|
||
|
|
def capture_gspro_window(window_title: str = "GSPro") -> Optional[np.ndarray]:
|
||
|
|
"""
|
||
|
|
Capture the GSPro window.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
window_title: GSPro window title
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Captured window as numpy array or None if not found
|
||
|
|
"""
|
||
|
|
return get_screen_capture().capture_window(window_title)
|
||
|
|
|
||
|
|
|
||
|
|
# Test function for development
|
||
|
|
def test_screen_capture():
|
||
|
|
"""Test screen capture functionality."""
|
||
|
|
print("Testing screen capture...")
|
||
|
|
|
||
|
|
capture = ScreenCapture()
|
||
|
|
|
||
|
|
# Get monitor information
|
||
|
|
monitors = capture.get_monitors()
|
||
|
|
print(f"Found {len(monitors)} monitors:")
|
||
|
|
for monitor in monitors:
|
||
|
|
print(
|
||
|
|
f" Monitor {monitor['index']}: {monitor['width']}x{monitor['height']} at ({monitor['left']}, {monitor['top']})"
|
||
|
|
)
|
||
|
|
|
||
|
|
# Capture primary screen
|
||
|
|
try:
|
||
|
|
screen = capture.capture_screen()
|
||
|
|
print(f"✓ Captured primary screen: {screen.shape}")
|
||
|
|
except Exception as e:
|
||
|
|
print(f"✗ Failed to capture screen: {e}")
|
||
|
|
|
||
|
|
# Test region capture
|
||
|
|
try:
|
||
|
|
region = capture.capture_region(100, 100, 640, 480)
|
||
|
|
print(f"✓ Captured region: {region.shape}")
|
||
|
|
except Exception as e:
|
||
|
|
print(f"✗ Failed to capture region: {e}")
|
||
|
|
|
||
|
|
# Test image to base64 conversion
|
||
|
|
try:
|
||
|
|
base64_str = capture.image_to_base64(region)
|
||
|
|
print(f"✓ Converted to base64: {len(base64_str)} chars")
|
||
|
|
except Exception as e:
|
||
|
|
print(f"✗ Failed to convert to base64: {e}")
|
||
|
|
|
||
|
|
# Test resolution presets
|
||
|
|
presets = ["480p", "720p", "1080p"]
|
||
|
|
for preset in presets:
|
||
|
|
width, height = capture.get_resolution_preset(preset)
|
||
|
|
print(f" {preset}: {width}x{height}")
|
||
|
|
|
||
|
|
capture.close()
|
||
|
|
print("\nScreen capture test complete!")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
test_screen_capture()
|