modules/element_processing.py

import io
import os
import time
import logging
import numpy as np
from typing import List, Dict, Tuple, Any, Union
from PIL import Image
from fastapi import BackgroundTasks
from modules.element_detector import detect_elements
from modules.element_annotation import annotate_image
from modules.config import ANNOTATION_DIR

logger = logging.getLogger(__name__)

# Unambiguous character set for code generation
LETTER_SET = "ACDEFHJKLMNPQRTUVWXY"
NUMBER_SET = "3479"

def generate_unique_codes(count: int) -> List[str]:
    """Generate unique two-character alphanumeric codes using visually unambiguous characters."""
    if count <= 0:
        return []
        
    codes = []
    letters = list(LETTER_SET)
    numbers = list(NUMBER_SET)
    
    # First use letter+number combinations
    for letter in letters:
        for number in numbers:
            codes.append(f"{letter}{number}")
            if len(codes) >= count:
                return codes
    
    # Then use letter+letter combinations if needed
    for first_letter in letters:
        for second_letter in letters:
            codes.append(f"{first_letter}{second_letter}")
            if len(codes) >= count:
                return codes
    
    return codes[:count]

def resolve_overlaps(text_elements: List[Dict], object_elements: List[Dict]) -> List[Dict]:
    """Resolve overlaps between text and object detection elements, prioritizing text."""
    # Safeguard against None values
    text_elements = text_elements or []
    object_elements = object_elements or []
        
    # Combine all elements
    all_elements = text_elements.copy()
    
    # Helper function to calculate IoU
    def calculate_iou(box1, box2):
        x1_min, y1_min, x1_max, y1_max = box1
        x2_min, y2_min, x2_max, y2_max = box2
        
        # Calculate intersection area
        x_left = max(x1_min, x2_min)
        y_top = max(y1_min, y2_min)
        x_right = min(x1_max, x2_max)
        y_bottom = min(y1_max, y2_max)
        
        if x_right < x_left or y_bottom < y_top:
            return 0.0
            
        intersection_area = (x_right - x_left) * (y_bottom - y_top)
        
        # Calculate union area
        box1_area = (x1_max - x1_min) * (y1_max - y1_min)
        box2_area = (x2_max - x2_min) * (y2_max - y2_min)
        union_area = box1_area + box2_area - intersection_area
        
        if union_area == 0:
            return 0.0
            
        return intersection_area / union_area
    
    # Check each object element against text elements
    for obj in object_elements:
        if "bbox" not in obj or len(obj["bbox"]) != 4:
            continue
            
        # Flag to check if object overlaps with any text element
        overlap = False
        obj_box = obj["bbox"]
        
        for text in text_elements:
            if "bbox" not in text or len(text["bbox"]) != 4:
                continue
                
            text_box = text["bbox"]
            iou = calculate_iou(obj_box, text_box)
            
            # If significant overlap, don't add the object element
            if iou > 0.5:
                overlap = True
                break
                
        if not overlap:
            all_elements.append(obj)
    
    return all_elements

def normalize_bounding_box(bbox: List[int], image_width: int, image_height: int) -> List[float]:
    """Convert absolute pixel coordinates to normalized coordinates (0.0-1.0)."""
    x_min, y_min, x_max, y_max = bbox
    
    # Ensure values are in valid range and convert any numpy types to Python native types
    x_min = float(max(0, min(int(x_min), image_width)))
    y_min = float(max(0, min(int(y_min), image_height)))
    x_max = float(max(0, min(int(x_max), image_width)))
    y_max = float(max(0, min(int(y_max), image_height)))
    
    return [
        x_min / float(image_width),
        y_min / float(image_height),
        x_max / float(image_width),
        y_max / float(image_height)
    ]

def calculate_center_coordinates(bbox: List[int], target_width: int = None, target_height: int = None, image_width: int = None, image_height: int = None) -> Tuple[int, int]:
    """
    Calculate the center coordinates of a bounding box as integers.
    If target dimensions are provided, scale the coordinates to match target screen dimensions.
    """
    x_min, y_min, x_max, y_max = bbox
    
    # Calculate center coordinates as floats first to maintain precision
    center_x_float = (x_min + x_max) / 2.0
    center_y_float = (y_min + y_max) / 2.0
    
    # If target dimensions are provided, scale the coordinates
    if all([target_width, target_height, image_width, image_height]):
        # Use floating point division for more accurate scaling
        x_scale = float(target_width) / float(image_width)
        y_scale = float(target_height) / float(image_height)
        
        center_x_float = center_x_float * x_scale
        center_y_float = center_y_float * y_scale
    
    # Only convert to integers at the final step to minimize rounding errors
    center_x = int(round(center_x_float))
    center_y = int(round(center_y_float))
        
    return center_x, center_y

def convert_numpy_types(obj: Any) -> Any:
    """Convert numpy types to Python native types to ensure JSON serialization works."""
    if isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, dict):
        return {k: convert_numpy_types(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [convert_numpy_types(i) for i in obj]
    else:
        return obj

async def process_screenshot(image_data: bytes, background_tasks: BackgroundTasks, screen_width: int = None, screen_height: int = None) -> Tuple[List[Dict], str]:
    """Process the screenshot to identify UI elements and assign unique codes."""
    start_time = time.time()
    
    try:
        # Convert image bytes to PIL Image
        image = Image.open(io.BytesIO(image_data))
        if image.mode == 'RGBA':
            image = image.convert('RGB')
        image_np = np.array(image)
        width, height = image.size
        
        # Detect elements (text and objects)
        text_elements, object_elements = await detect_elements(image_np)
        
        # Convert detection results to a common format
        elements = []
        
        # Process text elements
        for text in text_elements:
            if "text" in text and "bbox" in text:
                # Convert numpy types to Python native types
                bbox = [int(val) if isinstance(val, np.integer) else val for val in text["bbox"]]
                elements.append({
                    "type": "text",
                    "text_content": text["text"],
                    "bbox_pixels": bbox
                })
        
        # Process object elements
        for obj in object_elements:
            if "label" in obj and "bbox" in obj:
                # Convert numpy types to Python native types
                bbox = [int(val) if isinstance(val, np.integer) else val for val in obj["bbox"]]
                elements.append({
                    "type": "object",
                    "object_label": obj["label"],
                    "bbox_pixels": bbox
                })
        
        # Add normalized bounding boxes and center coordinates
        for element in elements:
            element["bbox_normalized"] = normalize_bounding_box(
                element["bbox_pixels"], width, height
            )
            # Calculate center coordinates, scaling to target screen dimensions if provided
            center_x, center_y = calculate_center_coordinates(
                element["bbox_pixels"], 
                target_width=screen_width, 
                target_height=screen_height,
                image_width=width,
                image_height=height
            )
            element["center_x"] = center_x
            element["center_y"] = center_y
        
        # Generate and assign unique codes
        if elements:
            codes = generate_unique_codes(len(elements))
            for i, element in enumerate(elements):
                element["code"] = codes[i]
        
        # Generate unique filename for the annotated image
        image_filename = f"annotated_{int(time.time())}.jpg"
        image_path = os.path.join(ANNOTATION_DIR, image_filename)
        
        # Process annotation (not in background anymore, we need the path)
        if elements:
            await annotate_image(
                image=image, 
                elements=elements, 
                filename=image_filename
            )
        
        logger.info(f"Element processing completed in {time.time() - start_time:.2f} seconds")
        logger.info(f"Found {len(elements)} elements: {len(text_elements)} text, {len(object_elements)} objects")
        
        # Convert any remaining numpy types to Python native types
        elements = convert_numpy_types(elements)
        
        return elements, image_path
        
    except Exception as e:
        logger.error(f"Error in process_screenshot: {str(e)}")
        return [], ""