Spaces:
Runtime error
Runtime error
File size: 3,364 Bytes
b9e2109 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | 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 [] |