Files
autoMate/gradio_ui/tools/screen_capture.py
2025-03-14 10:54:55 +08:00

23 lines
905 B
Python

from pathlib import Path
from uuid import uuid4
from PIL import Image
from .base import ToolError
from util import tool
OUTPUT_DIR = "./tmp/outputs"
def get_screenshot(resize: bool = False, target_width: int = 1920, target_height: int = 1080):
"""Capture screenshot by requesting from HTTP endpoint - returns native resolution unless resized"""
output_dir = Path(OUTPUT_DIR)
output_dir.mkdir(parents=True, exist_ok=True)
path = output_dir / f"screenshot_{uuid4().hex}.png"
try:
img_io = tool.capture_screen_with_cursor()
screenshot = Image.open(img_io)
if resize and screenshot.size != (target_width, target_height):
screenshot = screenshot.resize((target_width, target_height))
screenshot.save(path)
return screenshot, path
except Exception as e:
raise ToolError(f"Failed to capture screenshot: {str(e)}")