modules/element_detector.py

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