Spaces:
Paused
Paused
File size: 4,635 Bytes
f66643d | 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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | """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
|