import logging import asyncio import numpy as np from typing import List, Dict, Tuple, Optional from concurrent.futures import ThreadPoolExecutor from modules.detectors.base_detector import OCRBackend from modules.detectors.text_detector import EasyOCRBackend, GoogleVisionOCRBackend from modules.detectors.rectangular_detector import RectangleDetector from modules.config import OCR_BACKEND logger = logging.getLogger(__name__) # Global model instances ocr_backend = None ui_detector = None executor = ThreadPoolExecutor(max_workers=2) async def initialize_models(): """Initialize all detector models at application startup.""" global ocr_backend, ui_detector # Initialize OCR backend based on configuration ocr_backend_name = OCR_BACKEND.lower() if ocr_backend_name == "easyocr": ocr_backend = EasyOCRBackend() await ocr_backend.initialize() elif ocr_backend_name == "google": ocr_backend = GoogleVisionOCRBackend() await ocr_backend.initialize() else: raise ValueError(f"Unsupported OCR backend: {ocr_backend_name}") logger.info(f"OCR backend '{ocr_backend_name}' initialized") # Initialize UI element detector ui_detector = RectangleDetector() await ui_detector.initialize() logger.info("UI element detector initialized") async def detect_elements(image: np.ndarray) -> Tuple[List[Dict], List[Dict]]: """ Detect UI elements in the image using OCR and object detection. Args: image: Image as numpy array (OpenCV format) Returns: Tuple of (text_elements, object_elements) """ # Run OCR and object detection concurrently text_task = detect_text(image) object_task = detect_objects(image) text_elements, object_elements = await asyncio.gather(text_task, object_task) return text_elements, object_elements async def detect_text(image: np.ndarray) -> List[Dict]: """ Detect text elements in the image using the configured OCR backend. Args: image: Image as numpy array Returns: List of detected text elements """ if ocr_backend is None: logger.error("OCR backend not initialized") return [] try: # Run OCR detection in a separate thread to avoid blocking loop = asyncio.get_event_loop() text_elements = await loop.run_in_executor( executor, lambda: ocr_backend.detect(image) ) return text_elements except Exception as e: logger.error(f"Error detecting text: {str(e)}") return [] async def detect_objects(image: np.ndarray) -> List[Dict]: """ Detect UI objects in the image using object detection model. Args: image: Image as numpy array Returns: List of detected object elements """ if ui_detector is None: logger.error("UI detector not initialized") return [] try: # Run UI element detection in a separate thread loop = asyncio.get_event_loop() ui_elements = await loop.run_in_executor( executor, lambda: ui_detector.detect(image) ) return ui_elements except Exception as e: logger.error(f"Error detecting UI elements: {str(e)}") return []