OpenSpace/openspace/local_server/utils/screenshot.py

import platform
import os
import logging
from typing import Optional, Tuple
from PIL import Image
import pyautogui

logger = logging.getLogger(__name__)

platform_name = platform.system()


class ScreenshotHelper:
    def __init__(self):
        self.platform = platform_name
        self.adapter = None
        
        try:
            if platform_name == "Darwin":
                from ..platform_adapters.macos_adapter import MacOSAdapter
                self.adapter = MacOSAdapter()
            elif platform_name == "Linux":
                from ..platform_adapters.linux_adapter import LinuxAdapter
                self.adapter = LinuxAdapter()
            elif platform_name == "Windows":
                from ..platform_adapters.windows_adapter import WindowsAdapter
                self.adapter = WindowsAdapter()
        except ImportError as e:
            logger.warning(f"Failed to import platform adapter: {e}")
    
    def capture(self, output_path: str, with_cursor: bool = True) -> bool:
        try:
            # Ensure directory exists
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            if with_cursor and self.adapter:
                # Use platform-specific method to capture screenshot (with cursor)
                return self.adapter.capture_screenshot_with_cursor(output_path)
            else:
                # Use pyautogui to capture screenshot (without cursor)
                screenshot = pyautogui.screenshot()
                screenshot.save(output_path)
                logger.info(f"Screenshot successfully (without cursor): {output_path}")
                return True
                
        except Exception as e:
            logger.error(f"Screenshot failed: {e}")
            return False
    
    def capture_region(
        self, 
        output_path: str, 
        x: int, 
        y: int, 
        width: int, 
        height: int
    ) -> bool:
        """
        Capture specified screen region
        
        Args:
            output_path: Output path
            x: Starting x coordinate
            y: Starting y coordinate
            width: Width
            height: Height
            
        Returns:
            Whether successful
        """
        try:
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            screenshot = pyautogui.screenshot(region=(x, y, width, height))
            screenshot.save(output_path)
            logger.info(f"Region screenshot successfully: {output_path}")
            return True
            
        except Exception as e:
            logger.error(f"Region screenshot failed: {e}")
            return False
    
    def get_screen_size(self) -> Tuple[int, int]:
        """
        Get screen size
        
        Returns:
            (width, height)
        """
        try:
            size = pyautogui.size()
            return (size.width, size.height)
        except Exception as e:
            logger.error(f"Failed to get screen size: {e}")
            return (1920, 1080)  # Default value
    
    def get_cursor_position(self) -> Tuple[int, int]:
        """
        Get cursor position
        
        Returns:
            (x, y)
        """
        try:
            pos = pyautogui.position()
            return (pos.x, pos.y)
        except Exception as e:
            logger.error(f"Failed to get cursor position: {e}")
            return (0, 0)
    
    def capture_to_base64(self, with_cursor: bool = True) -> Optional[str]:
        """
        Capture screenshot and convert to base64
        
        Args:
            with_cursor: Whether to include cursor
            
        Returns:
            Base64 encoded image string
        """
        import tempfile
        import base64
        
        try:
            # Create temporary file
            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
                tmp_path = tmp.name
            
            # Capture screenshot
            if self.capture(tmp_path, with_cursor):
                # Read and encode
                with open(tmp_path, 'rb') as f:
                    img_data = f.read()
                    img_base64 = base64.b64encode(img_data).decode('utf-8')
                
                # Delete temporary file
                os.remove(tmp_path)
                
                return img_base64
            else:
                if os.path.exists(tmp_path):
                    os.remove(tmp_path)
                return None
                
        except Exception as e:
            logger.error(f"Failed to convert screenshot to base64: {e}")
            return None
    
    def compare_screenshots(self, path1: str, path2: str) -> float:
        """
        Compare similarity between two screenshots
        
        Args:
            path1: First image path
            path2: Second image path
            
        Returns:
            Similarity (0-1), 1 means identical
        """
        try:
            from PIL import ImageChops
            import math
            import operator
            from functools import reduce
            
            img1 = Image.open(path1)
            img2 = Image.open(path2)
            
            # Ensure same size
            if img1.size != img2.size:
                # Resize to same size
                img2 = img2.resize(img1.size)
            
            # Calculate difference
            diff = ImageChops.difference(img1, img2)
            
            # Calculate statistics
            stat = diff.histogram()
            sum_of_squares = reduce(
                operator.add,
                map(lambda h, i: h * (i ** 2), stat, range(len(stat)))
            )
            
            # Calculate RMS
            rms = math.sqrt(sum_of_squares / float(img1.size[0] * img1.size[1]))
            
            # Normalize to 0-1, RMS max value is approximately 441 (for RGB)
            similarity = 1 - (rms / 441.0)
            
            return max(0, min(1, similarity))
            
        except Exception as e:
            logger.error(f"Failed to compare screenshots: {e}")
            return 0.0
    
    def annotate_screenshot(
        self, 
        input_path: str, 
        output_path: str, 
        annotations: list
    ) -> bool:
        """
        Add annotations to screenshot
        
        Args:
            input_path: Input image path
            output_path: Output image path
            annotations: List of annotations, each annotation is a dict:
                        {'type': 'rectangle'/'text', 'x': int, 'y': int, 
                         'width': int, 'height': int, 'text': str, 'color': tuple}
            
        Returns:
            Whether successful
        """
        try:
            from PIL import ImageDraw, ImageFont
            
            img = Image.open(input_path)
            draw = ImageDraw.Draw(img)
            
            for annotation in annotations:
                ann_type = annotation.get('type', 'rectangle')
                color = annotation.get('color', (255, 0, 0))
                
                if ann_type == 'rectangle':
                    x = annotation.get('x', 0)
                    y = annotation.get('y', 0)
                    width = annotation.get('width', 100)
                    height = annotation.get('height', 100)
                    
                    draw.rectangle(
                        [(x, y), (x + width, y + height)],
                        outline=color,
                        width=2
                    )
                    
                elif ann_type == 'text':
                    x = annotation.get('x', 0)
                    y = annotation.get('y', 0)
                    text = annotation.get('text', '')
                    
                    try:
                        font = ImageFont.truetype("Arial.ttf", 20)
                    except:
                        font = ImageFont.load_default()
                    
                    draw.text((x, y), text, fill=color, font=font)
            
            img.save(output_path)
            logger.info(f"Annotated screenshot successfully: {output_path}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to annotate screenshot: {e}")
            return False