Source code for fmusvid.capture.screen

"""
Screen capture functionality for FMUS-VID.

This module provides tools for capturing screenshots of the screen or specific regions.
"""

import os
import platform
from pathlib import Path
from typing import Optional, Tuple, Union, List
import tempfile
import time

import numpy as np
from PIL import Image, ImageGrab
import cv2  # Added for image processing

# Define a simple Frame class for the capture module instead of importing from core
[docs] class Frame: """Simplified Frame class for captured screenshots."""
[docs] def __init__(self, frame_data): """Initialize a Frame with numpy array data.""" self.data = frame_data
[docs] def save(self, path): """Save the frame to a file.""" path = Path(path) cv2.imwrite(str(path), cv2.cvtColor(self.data, cv2.COLOR_RGB2BGR)) return path
[docs] def to_numpy(self): """Return the frame as a numpy array.""" return self.data
[docs] def to_pil(self): """Convert to PIL Image.""" return Image.fromarray(self.data)
[docs] def capture_screen(output_path: Optional[Union[str, Path]] = None) -> Union[Frame, Path]: """ Capture a screenshot of the entire screen. Args: output_path: Optional path to save the screenshot. If None, returns a Frame object. Returns: Frame object if output_path is None, otherwise the Path to the saved image. Example: >>> # Capture and get Frame object >>> frame = fmusvid.capture.capture_screen() >>> # Capture and save to file >>> path = fmusvid.capture.capture_screen("screenshot.png") """ # Capture the screen screenshot = ImageGrab.grab() # Convert to numpy array for processing frame_array = np.array(screenshot) # Create a Frame object frame = Frame(frame_array) # If output path is provided, save the image if output_path: output_path = Path(output_path) return frame.save(output_path) return frame
[docs] def capture_region(region: Tuple[int, int, int, int], output_path: Optional[Union[str, Path]] = None) -> Union[Frame, Path]: """ Capture a screenshot of a specific region of the screen. Args: region: Tuple of (left, top, right, bottom) coordinates output_path: Optional path to save the screenshot. If None, returns a Frame object. Returns: Frame object if output_path is None, otherwise the Path to the saved image. Example: >>> # Capture region (100, 100, 500, 400) and get Frame object >>> frame = fmusvid.capture.capture_region((100, 100, 500, 400)) >>> # Capture region and save to file >>> path = fmusvid.capture.capture_region((100, 100, 500, 400), "region.png") """ # Validate region coordinates left, top, right, bottom = region if left >= right or top >= bottom: raise ValueError("Invalid region coordinates. Ensure left < right and top < bottom.") # Capture the region screenshot = ImageGrab.grab(bbox=region) # Convert to numpy array for processing frame_array = np.array(screenshot) # Create a Frame object frame = Frame(frame_array) # If output path is provided, save the image if output_path: output_path = Path(output_path) return frame.save(output_path) return frame
[docs] def select_region() -> Tuple[int, int, int, int]: """ Interactive tool to select a region of the screen. Returns: Tuple of (left, top, right, bottom) coordinates Example: >>> # Select a region interactively >>> region = fmusvid.capture.select_region() >>> # Capture the selected region >>> frame = fmusvid.capture.capture_region(region) """ try: import tkinter as tk from tkinter import simpledialog except ImportError: raise ImportError("Tkinter is required for interactive region selection. " "Install it or specify region coordinates manually.") # Create a full screen transparent window root = tk.Tk() root.attributes('-alpha', 0.3) root.attributes('-fullscreen', True) # Variables to store the selection coordinates start_x = start_y = end_x = end_y = 0 selection_made = False # Create a canvas canvas = tk.Canvas(root, cursor="cross") canvas.pack(fill=tk.BOTH, expand=True) # Rectangle ID for the selection rect_id = None def on_mouse_down(event): nonlocal start_x, start_y, rect_id start_x, start_y = event.x, event.y # Create a rectangle if it doesn't exist if rect_id: canvas.delete(rect_id) rect_id = canvas.create_rectangle(start_x, start_y, start_x, start_y, outline='red', width=2) def on_mouse_move(event): nonlocal rect_id if rect_id: canvas.coords(rect_id, start_x, start_y, event.x, event.y) def on_mouse_up(event): nonlocal end_x, end_y, selection_made end_x, end_y = event.x, event.y selection_made = True root.quit() # Bind mouse events canvas.bind("<ButtonPress-1>", on_mouse_down) canvas.bind("<B1-Motion>", on_mouse_move) canvas.bind("<ButtonRelease-1>", on_mouse_up) # Run the main loop root.mainloop() # Clean up root.destroy() # Sort coordinates so left < right and top < bottom left = min(start_x, end_x) top = min(start_y, end_y) right = max(start_x, end_x) bottom = max(start_y, end_y) return (left, top, right, bottom)