"""
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)