Spaces:
Paused
Paused
| """OCR model: text detection + recognition""" | |
| from __future__ import annotations | |
| import logging | |
| from typing import Any, Optional, Tuple | |
| from PIL import Image | |
| from pdf2zh.parser.ai_models.base import BaseImageToTextModel | |
| logger = logging.getLogger(__name__) | |
| class SuryaOCRModel(BaseImageToTextModel): | |
| """ | |
| Wraps Surya's DetectionPredictor + RecognitionPredictor. | |
| Models are loaded lazily upon first inference call. | |
| """ | |
| model_name = "SuryaOCR" | |
| def __init__( | |
| self, | |
| detector_blank_threshold: Optional[float] = None, | |
| detector_text_threshold: Optional[float] = None, | |
| ) -> None: | |
| super().__init__() | |
| self.detector_blank_threshold = detector_blank_threshold | |
| self.detector_text_threshold = detector_text_threshold | |
| self.foundation_predictor: Any = None | |
| self.detection_predictor: Any = None | |
| self.recognition_predictor: Any = None | |
| def load_model(self) -> None: | |
| logger.info( | |
| "Initializing %s and loading models into memory...", self.model_name | |
| ) | |
| from surya.detection import DetectionPredictor | |
| from surya.foundation import FoundationPredictor | |
| from surya.recognition import RecognitionPredictor | |
| from surya.settings import settings | |
| if self.detector_text_threshold is not None: | |
| settings.DETECTOR_TEXT_THRESHOLD = self.detector_text_threshold | |
| if self.detector_blank_threshold is not None: | |
| settings.DETECTOR_BLANK_THRESHOLD = self.detector_blank_threshold | |
| self.foundation_predictor = FoundationPredictor() | |
| logger.info("Loaded FoundationPredictor (OCR backbone)") | |
| self.detection_predictor = DetectionPredictor() | |
| logger.info("Loaded DetectionPredictor") | |
| self.recognition_predictor = RecognitionPredictor(self.foundation_predictor) | |
| logger.info("Loaded RecognitionPredictor") | |
| self.model = self.recognition_predictor | |
| def unload_model(self) -> None: | |
| if self.model is not None: | |
| import torch | |
| logger.info("Unloading all %s predictors from VRAM...", self.model_name) | |
| del self.foundation_predictor | |
| del self.detection_predictor | |
| del self.recognition_predictor | |
| del self.model | |
| self.foundation_predictor = None | |
| self.detection_predictor = None | |
| self.recognition_predictor = None | |
| self.model = None | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| def prepare( | |
| self, | |
| images: list[Image.Image], | |
| highres_images: list[Image.Image] | None = None, | |
| *args: Any, | |
| **kwargs: Any, | |
| ) -> Tuple[list[Image.Image], list[Image.Image] | None]: | |
| """ | |
| Preprocess raw images before inference. | |
| """ | |
| return images, highres_images | |
| def predict( | |
| self, | |
| prepared_inputs: Tuple[list[Image.Image], list[Image.Image] | None], | |
| *args: Any, | |
| math_mode: bool = False, | |
| task_names: list[Any] | None = None, | |
| bboxes: list[Any] | None = None, | |
| detection_batch_size: int | None = None, | |
| ocr_batch_size: int | None = None, | |
| **kwargs: Any, | |
| ) -> list[Any]: | |
| """ | |
| Run full-page OCR (detection -> recognition) on prepared images. | |
| """ | |
| images, highres_images = prepared_inputs | |
| run_kwargs: dict[str, Any] = {"math_mode": True, "return_words": False} | |
| if not math_mode: | |
| logger.info("Running OCR with detection + recognition") | |
| run_kwargs.update( | |
| { | |
| "det_predictor": self.detection_predictor, | |
| "detection_batch_size": detection_batch_size, | |
| "recognition_batch_size": ocr_batch_size, | |
| "highres_images": highres_images, | |
| } | |
| ) | |
| else: | |
| logger.info("Running OCR in math mode (LaTeX recognition)") | |
| run_kwargs.update( | |
| { | |
| "recognition_batch_size": ocr_batch_size, | |
| } | |
| ) | |
| if task_names is not None: | |
| run_kwargs["task_names"] = task_names | |
| if bboxes is not None: | |
| run_kwargs["bboxes"] = bboxes | |
| raw_results = self.recognition_predictor(images, **run_kwargs) | |
| return raw_results | |
| def postprocess( | |
| self, raw_results: list[Any], *args: Any, **kwargs: Any | |
| ) -> list[Any]: | |
| """ | |
| Format raw Surya outputs into the final desired structure. | |
| """ | |
| return raw_results | |