Spaces:
Running on Zero
Running on Zero
File size: 4,551 Bytes
b411c37 819da4b b411c37 819da4b b411c37 | 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 | from __future__ import annotations
import threading
from dataclasses import dataclass
from typing import Any
from PIL import Image
from ocr_studio.config import (
MODE_PRECISE,
MODEL_ID,
MODEL_REVISION,
OCR_MAX_PIXELS,
SPOTTING_MAX_PIXELS,
SPOTTING_UPSCALE_THRESHOLD,
MAX_NEW_TOKENS,
)
from ocr_studio.spotting import TextSpan, parse_spans, spans_to_text, strip_special_tokens
PROMPTS = {
"ocr": "OCR:",
"spotting": "Spotting:",
}
@dataclass
class InferenceResult:
raw_text: str
display_text: str
spans: list[TextSpan]
task: str
class PaddleOcrVlEngine:
def __init__(self) -> None:
self.model: Any = None
self.processor: Any = None
self.device: Any = None
self._lock = threading.Lock()
def load(self) -> None:
if self.model is not None:
return
import torch
from transformers import AutoConfig, AutoModelForImageTextToText, AutoProcessor
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
config = AutoConfig.from_pretrained(MODEL_ID, revision=MODEL_REVISION)
if not hasattr(config, "text_config") and hasattr(config, "get_text_config"):
config.text_config = config.get_text_config()
processor = AutoProcessor.from_pretrained(
MODEL_ID,
revision=MODEL_REVISION,
trust_remote_code=False,
)
model = AutoModelForImageTextToText.from_pretrained(
MODEL_ID,
config=config,
revision=MODEL_REVISION,
torch_dtype=dtype,
trust_remote_code=False,
low_cpu_mem_usage=True,
)
model = model.to(device).eval()
self.model = model
self.processor = processor
self.device = device
def _prepare_image(self, image: Image.Image, task: str) -> Image.Image:
prepared = image.convert("RGB")
if (
task == "spotting"
and prepared.width < SPOTTING_UPSCALE_THRESHOLD
and prepared.height < SPOTTING_UPSCALE_THRESHOLD
):
prepared = prepared.resize(
(prepared.width * 2, prepared.height * 2),
Image.Resampling.LANCZOS,
)
return prepared
def recognize(self, image: Image.Image, mode: str) -> InferenceResult:
self.load()
import torch
task = "spotting" if mode == MODE_PRECISE else "ocr"
work_image = self._prepare_image(image, task)
max_pixels = SPOTTING_MAX_PIXELS if task == "spotting" else OCR_MAX_PIXELS
messages = [
{
"role": "user",
"content": [
{"type": "image", "image": work_image},
{"type": "text", "text": PROMPTS[task]},
],
}
]
image_processor = self.processor.image_processor
min_pixels = getattr(image_processor, "min_pixels", None)
if min_pixels is None:
size_cfg = getattr(image_processor, "size", {}) or {}
min_pixels = size_cfg.get("shortest_edge") or size_cfg.get("min_pixels") or (16 * 28 * 28)
inputs = self.processor.apply_chat_template(
messages,
add_generation_prompt=True,
tokenize=True,
return_dict=True,
return_tensors="pt",
images_kwargs={
"size": {
"shortest_edge": int(min_pixels),
"longest_edge": max_pixels,
}
},
)
inputs = inputs.to(self.model.device)
with self._lock, torch.inference_mode():
generated = self.model.generate(
**inputs,
max_new_tokens=MAX_NEW_TOKENS,
do_sample=False,
)
prompt_len = inputs["input_ids"].shape[-1]
decoded = self.processor.decode(
generated[0][prompt_len:],
skip_special_tokens=True,
clean_up_tokenization_spaces=False,
).strip()
spans = parse_spans(decoded, image.width, image.height)
display = spans_to_text(spans, strip_special_tokens(decoded)).strip()
return InferenceResult(raw_text=decoded, display_text=display, spans=spans, task=task)
|