220 lines
No EOL
7 KiB
Python
220 lines
No EOL
7 KiB
Python
"""Screen capture utilities for taking game screenshots.
|
|
|
|
Provides efficient screenshot capture using multiple backends (mss, PIL)
|
|
with support for specific regions and window targeting.
|
|
"""
|
|
|
|
from typing import Tuple, Optional, Dict, Any
|
|
from dataclasses import dataclass
|
|
import time
|
|
import logging
|
|
|
|
import numpy as np
|
|
from PIL import Image, ImageGrab
|
|
import mss
|
|
import cv2
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class ScreenRegion:
|
|
"""Defines a rectangular region of the screen to capture."""
|
|
|
|
x: int
|
|
y: int
|
|
width: int
|
|
height: int
|
|
|
|
@property
|
|
def bounds(self) -> Tuple[int, int, int, int]:
|
|
"""Return region as (left, top, right, bottom) tuple."""
|
|
return (self.x, self.y, self.x + self.width, self.y + self.height)
|
|
|
|
@property
|
|
def mss_bounds(self) -> Dict[str, int]:
|
|
"""Return region in MSS format."""
|
|
return {
|
|
"top": self.y,
|
|
"left": self.x,
|
|
"width": self.width,
|
|
"height": self.height,
|
|
}
|
|
|
|
|
|
class ScreenCapture:
|
|
"""High-performance screen capture with multiple backends."""
|
|
|
|
def __init__(self, backend: str = "mss", monitor: int = 1):
|
|
"""Initialize screen capture.
|
|
|
|
Args:
|
|
backend: Capture backend ("mss" or "pil")
|
|
monitor: Monitor number to capture from (1-indexed)
|
|
"""
|
|
self.backend = backend
|
|
self.monitor = monitor
|
|
self._mss_instance: Optional[mss.mss] = None
|
|
self._monitor_info: Optional[Dict[str, int]] = None
|
|
|
|
if backend == "mss":
|
|
self._initialize_mss()
|
|
|
|
def _initialize_mss(self) -> None:
|
|
"""Initialize MSS backend."""
|
|
try:
|
|
self._mss_instance = mss.mss()
|
|
monitors = self._mss_instance.monitors
|
|
|
|
if self.monitor >= len(monitors):
|
|
logger.warning(f"Monitor {self.monitor} not found, using primary")
|
|
self.monitor = 1
|
|
|
|
self._monitor_info = monitors[self.monitor]
|
|
logger.info(f"Initialized MSS capture for monitor {self.monitor}: "
|
|
f"{self._monitor_info['width']}x{self._monitor_info['height']}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize MSS: {e}")
|
|
self.backend = "pil"
|
|
|
|
def capture_screen(self, region: Optional[ScreenRegion] = None) -> np.ndarray:
|
|
"""Capture screenshot of screen or region.
|
|
|
|
Args:
|
|
region: Specific region to capture, or None for full screen
|
|
|
|
Returns:
|
|
Screenshot as numpy array in BGR format (for OpenCV compatibility)
|
|
"""
|
|
try:
|
|
if self.backend == "mss":
|
|
return self._capture_mss(region)
|
|
else:
|
|
return self._capture_pil(region)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Screen capture failed: {e}")
|
|
# Fallback to empty image
|
|
return np.zeros((100, 100, 3), dtype=np.uint8)
|
|
|
|
def _capture_mss(self, region: Optional[ScreenRegion]) -> np.ndarray:
|
|
"""Capture using MSS backend."""
|
|
if not self._mss_instance:
|
|
raise RuntimeError("MSS not initialized")
|
|
|
|
if region:
|
|
monitor = region.mss_bounds
|
|
else:
|
|
monitor = self._monitor_info or self._mss_instance.monitors[self.monitor]
|
|
|
|
# MSS returns BGRA format
|
|
screenshot = self._mss_instance.grab(monitor)
|
|
img_array = np.frombuffer(screenshot.rgb, dtype=np.uint8)
|
|
img_array = img_array.reshape((screenshot.height, screenshot.width, 3))
|
|
|
|
# Convert RGB to BGR for OpenCV
|
|
return cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
|
|
|
|
def _capture_pil(self, region: Optional[ScreenRegion]) -> np.ndarray:
|
|
"""Capture using PIL backend."""
|
|
if region:
|
|
bbox = region.bounds
|
|
else:
|
|
bbox = None
|
|
|
|
# PIL returns RGB format
|
|
screenshot = ImageGrab.grab(bbox=bbox)
|
|
img_array = np.array(screenshot)
|
|
|
|
# Convert RGB to BGR for OpenCV
|
|
return cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
|
|
|
|
def save_screenshot(self, filename: str, region: Optional[ScreenRegion] = None) -> bool:
|
|
"""Save screenshot to file.
|
|
|
|
Args:
|
|
filename: Output filename
|
|
region: Region to capture, or None for full screen
|
|
|
|
Returns:
|
|
True if successful, False otherwise
|
|
"""
|
|
try:
|
|
img = self.capture_screen(region)
|
|
return cv2.imwrite(filename, img)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to save screenshot: {e}")
|
|
return False
|
|
|
|
def get_screen_size(self) -> Tuple[int, int]:
|
|
"""Get screen dimensions.
|
|
|
|
Returns:
|
|
(width, height) tuple
|
|
"""
|
|
if self.backend == "mss" and self._monitor_info:
|
|
return (self._monitor_info["width"], self._monitor_info["height"])
|
|
else:
|
|
# Use PIL as fallback
|
|
screenshot = ImageGrab.grab()
|
|
return screenshot.size
|
|
|
|
def find_window(self, window_title: str) -> Optional[ScreenRegion]:
|
|
"""Find window by title and return its region.
|
|
|
|
Args:
|
|
window_title: Partial or full window title to search for
|
|
|
|
Returns:
|
|
ScreenRegion if window found, None otherwise
|
|
|
|
Note:
|
|
This is a placeholder - actual implementation would use
|
|
platform-specific window enumeration (e.g., Windows API, X11)
|
|
"""
|
|
# TODO: Implement window finding
|
|
logger.warning("Window finding not implemented yet")
|
|
return None
|
|
|
|
def benchmark_capture(self, iterations: int = 100) -> Dict[str, float]:
|
|
"""Benchmark capture performance.
|
|
|
|
Args:
|
|
iterations: Number of captures to perform
|
|
|
|
Returns:
|
|
Performance statistics
|
|
"""
|
|
logger.info(f"Benchmarking {self.backend} backend ({iterations} iterations)")
|
|
|
|
start_time = time.perf_counter()
|
|
|
|
for _ in range(iterations):
|
|
self.capture_screen()
|
|
|
|
end_time = time.perf_counter()
|
|
total_time = end_time - start_time
|
|
avg_time = total_time / iterations
|
|
fps = iterations / total_time
|
|
|
|
stats = {
|
|
"backend": self.backend,
|
|
"iterations": iterations,
|
|
"total_time": total_time,
|
|
"avg_time_ms": avg_time * 1000,
|
|
"fps": fps,
|
|
}
|
|
|
|
logger.info(f"Benchmark results: {avg_time*1000:.2f}ms avg, {fps:.1f} FPS")
|
|
return stats
|
|
|
|
def __enter__(self):
|
|
"""Context manager entry."""
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Context manager exit."""
|
|
if self._mss_instance:
|
|
self._mss_instance.close() |