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 [], ""