Spaces:
Runtime error
Runtime error
| import os | |
| import logging | |
| import numpy as np | |
| from typing import List, Dict, Optional | |
| from modules.detectors.base_detector import OCRBackend | |
| from modules.config import OCR_NUM_WORKERS, OCR_CONFIDENCE_THRESHOLD, GOOGLE_APPLICATION_CREDENTIALS | |
| from modules.config import PYTORCH_ENABLE_MPS_FALLBACK | |
| logger = logging.getLogger(__name__) | |
| # Singleton pattern for EasyOCR reader | |
| class EasyOCRSingleton: | |
| _instance: Optional['EasyOCRSingleton'] = None | |
| reader = None | |
| device = "cpu" | |
| def get_instance(cls): | |
| if cls._instance is None: | |
| cls._instance = cls() | |
| return cls._instance | |
| async def initialize(cls, languages=["en"], gpu=True, num_workers=None): | |
| """Initialize the EasyOCR singleton instance.""" | |
| instance = cls.get_instance() | |
| if instance.reader is not None: | |
| logger.info("EasyOCR reader already initialized, reusing instance") | |
| return instance | |
| try: | |
| import easyocr | |
| import torch | |
| # Check for available hardware acceleration | |
| instance.device = "cpu" | |
| if torch.cuda.is_available(): | |
| logger.info("CUDA is available, using GPU acceleration") | |
| gpu = True | |
| instance.device = "cuda" | |
| elif hasattr(torch, 'mps') and torch.mps.is_available(): | |
| logger.info("MPS is available, using Apple Silicon acceleration") | |
| gpu = True | |
| instance.device = "mps" | |
| else: | |
| logger.warning("No hardware acceleration available, falling back to CPU") | |
| gpu = False | |
| # Get num_workers from config or use provided value | |
| if num_workers is None: | |
| num_workers = OCR_NUM_WORKERS | |
| # Initialize EasyOCR reader with appropriate device | |
| # EasyOCR doesn't accept device parameter directly, for MPS we need to handle differently | |
| if instance.device == "mps": | |
| # For MPS, we need to set the device in torch before initializing | |
| os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = PYTORCH_ENABLE_MPS_FALLBACK | |
| torch.set_default_device(instance.device) | |
| reader_args = { | |
| "lang_list": languages, | |
| "gpu": gpu | |
| } | |
| instance.reader = easyocr.Reader(**reader_args) | |
| logger.info(f"EasyOCR initialized with device={instance.device}, num_workers={num_workers}") | |
| return instance | |
| except Exception as e: | |
| logger.error(f"Failed to initialize EasyOCR: {str(e)}") | |
| raise | |
| class EasyOCRBackend(OCRBackend): | |
| """EasyOCR implementation for text detection.""" | |
| def __init__(self): | |
| self.languages = ["en"] | |
| self.gpu = True # Use GPU if available | |
| self.num_workers = OCR_NUM_WORKERS | |
| async def initialize(self): | |
| """Initialize the EasyOCR model using the singleton pattern.""" | |
| try: | |
| # Initialize the singleton instance | |
| await EasyOCRSingleton.initialize( | |
| languages=self.languages, | |
| gpu=self.gpu, | |
| num_workers=self.num_workers | |
| ) | |
| logger.info(f"EasyOCR backend initialized with num_workers={self.num_workers}") | |
| except Exception as e: | |
| logger.error(f"Failed to initialize EasyOCR: {str(e)}") | |
| raise | |
| def detect(self, image: np.ndarray) -> List[Dict]: | |
| """Detect text in an image using EasyOCR.""" | |
| # Get the reader from singleton | |
| ocr_instance = EasyOCRSingleton.get_instance() | |
| if ocr_instance.reader is None: | |
| logger.error("EasyOCR model not initialized") | |
| return [] | |
| try: | |
| # Detect text using the shared singleton reader | |
| results = ocr_instance.reader.readtext(image) | |
| # Get confidence threshold from config | |
| confidence_threshold = OCR_CONFIDENCE_THRESHOLD | |
| # Process and format results | |
| text_elements = [] | |
| for bbox, text, confidence in results: | |
| if confidence < confidence_threshold: # Filter out low confidence results | |
| continue | |
| # Convert EasyOCR points format to [x_min, y_min, x_max, y_max] | |
| points = np.array(bbox) | |
| x_min = int(min(points[:, 0])) | |
| y_min = int(min(points[:, 1])) | |
| x_max = int(max(points[:, 0])) | |
| y_max = int(max(points[:, 1])) | |
| text_elements.append({ | |
| "text": text, | |
| "confidence": float(confidence), | |
| "bbox": [x_min, y_min, x_max, y_max] | |
| }) | |
| return text_elements | |
| except Exception as e: | |
| logger.error(f"Error in EasyOCR detection: {str(e)}") | |
| return [] | |
| class GoogleVisionOCRBackend(OCRBackend): | |
| """Google Cloud Vision API implementation for text detection.""" | |
| def __init__(self): | |
| self.client = None | |
| self.credentials_path = GOOGLE_APPLICATION_CREDENTIALS | |
| async def initialize(self): | |
| """Initialize the Google Cloud Vision client.""" | |
| try: | |
| from google.cloud import vision | |
| if not self.credentials_path: | |
| raise ValueError("GOOGLE_APPLICATION_CREDENTIALS environment variable not set") | |
| if not os.path.exists(self.credentials_path): | |
| raise FileNotFoundError(f"Google Cloud credentials file not found: {self.credentials_path}") | |
| # Initialize Vision client | |
| self.client = vision.ImageAnnotatorClient() | |
| logger.info("Google Cloud Vision client initialized") | |
| except Exception as e: | |
| logger.error(f"Failed to initialize Google Vision: {str(e)}") | |
| raise | |
| def detect(self, image: np.ndarray) -> List[Dict]: | |
| """Detect text in an image using Google Cloud Vision API.""" | |
| if self.client is None: | |
| logger.error("Google Cloud Vision client not initialized") | |
| return [] | |
| try: | |
| from google.cloud import vision | |
| import cv2 | |
| # Convert numpy array to bytes | |
| _, buffer = cv2.imencode(".jpg", image) | |
| content = buffer.tobytes() | |
| # Create vision image | |
| vision_image = vision.Image(content=content) | |
| # Detect text | |
| response = self.client.text_detection(image=vision_image) | |
| if response.error.message: | |
| logger.error(f"Google Vision API error: {response.error.message}") | |
| return [] | |
| # Process and format results | |
| text_elements = [] | |
| for text_annotation in response.text_annotations[1:]: # Skip the first annotation (full page text) | |
| vertices = text_annotation.bounding_poly.vertices | |
| # Convert Vision API vertices to [x_min, y_min, x_max, y_max] | |
| x_coords = [vertex.x for vertex in vertices] | |
| y_coords = [vertex.y for vertex in vertices] | |
| x_min = min(x_coords) | |
| y_min = min(y_coords) | |
| x_max = max(x_coords) | |
| y_max = max(y_coords) | |
| text_elements.append({ | |
| "text": text_annotation.description, | |
| "bbox": [x_min, y_min, x_max, y_max] | |
| }) | |
| return text_elements | |
| except Exception as e: | |
| logger.error(f"Error in Google Vision detection: {str(e)}") | |
| return [] |