Spaces:
Running on Zero
Running on Zero
| """Azerbaijani line & page OCR — CRNN + CTC recogniser (LocalDoc/azerbaijani-ocr-crnn). | |
| Two modes: | |
| * Line — the recogniser alone on a cropped line (the model's native input). | |
| * Page — PP-OCRv5 detection finds the lines, the recogniser reads them. | |
| Vendored from the model repo (CC-BY-4.0): `modeling.py` and the detection / | |
| reading-order logic of `pipeline.py`. | |
| """ | |
| import os | |
| # Must precede the torch import: allocator fragmentation under the transient | |
| # pressure of many wide crops (see CUDACachingAllocator NVML assert). | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| os.environ.setdefault("FLAGS_use_mkldnn", "0") # oneDNN clashes with the PIR executor | |
| os.environ.setdefault("GLOG_minloglevel", "2") | |
| import spaces # noqa: F401 — must precede any CUDA-touching import | |
| import tempfile | |
| import time | |
| import threading | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from huggingface_hub import hf_hub_download | |
| import gradio as gr | |
| from modeling import load, preprocess, ctc_greedy, HEIGHT, DOWNSAMPLE | |
| MODEL_ID = "LocalDoc/azerbaijani-ocr-crnn" | |
| DET_MODEL = "PP-OCRv5_server_det" | |
| UNCLIP = 1.5 # box inflation; raise if ç ş ğ descenders get clipped | |
| PAD_FRAC = 0.12 # vertical headroom for descenders | |
| MAX_WIDTH = 1200 | |
| BATCH = 12 | |
| # --------------------------------------------------------------------------- model | |
| _path = hf_hub_download(MODEL_ID, "model.pt") | |
| _model, _itos = load(_path, "cuda") # eager .to("cuda"); ZeroGPU packs it at startup | |
| _model.eval() | |
| print(f"recogniser ready (training={_model.training})", flush=True) | |
| # ------------------------------------------------------------------- detector (lazy) | |
| _det = None | |
| _det_lock = threading.Lock() | |
| def _detector(): | |
| """PP-OCRv5 detector, built once and reused by the warm worker. | |
| Constructed lazily: `import spaces` has already run by the time this is | |
| first called, so importing paddle inside a `@spaces.GPU` call is safe and | |
| paddle never becomes a hard module-scope dependency. | |
| """ | |
| global _det | |
| if _det is None: | |
| with _det_lock: | |
| if _det is None: | |
| from paddleocr import TextDetection | |
| _det = TextDetection(model_name=DET_MODEL, enable_mkldnn=False, | |
| unclip_ratio=UNCLIP) | |
| return _det | |
| # Warm the detector's weights into the image cache at build/startup time, so a | |
| # cold GPU worker doesn't spend its first request downloading them. Paddle is a | |
| # CPU build, so this does not initialise CUDA. Non-fatal if it fails: page mode | |
| # then pays the download once on its first call. | |
| try: | |
| _detector() | |
| print("PP-OCRv5 detector ready", flush=True) | |
| except Exception as _e: # noqa: BLE001 | |
| print(f"detector pre-warm skipped ({_e!r}); will build lazily", flush=True) | |
| # ------------------------------------------------------- reading order (from pipeline.py) | |
| def _reading_order(boxes): | |
| """Cluster boxes into rows by y, then left to right within each row.""" | |
| if not boxes: | |
| return [] | |
| arr = [np.array(b) for b in boxes] | |
| tol = max(6.0, float(np.median([np.ptp(b[:, 1]) for b in arr])) * 0.6) | |
| idx = sorted(range(len(arr)), key=lambda i: arr[i][:, 1].mean()) | |
| rows, cur = [], [idx[0]] | |
| for i in idx[1:]: | |
| if abs(arr[i][:, 1].mean() - arr[cur[-1]][:, 1].mean()) <= tol: | |
| cur.append(i) | |
| else: | |
| rows.append(cur) | |
| cur = [i] | |
| rows.append(cur) | |
| out = [] | |
| for row in rows: | |
| out += sorted(row, key=lambda i: arr[i][:, 0].min()) | |
| return [boxes[i] for i in out] | |
| def _pad(boxes, height): | |
| """Pad above and below, never more than half the gap to the neighbouring line.""" | |
| order = sorted(range(len(boxes)), key=lambda i: boxes[i][1]) | |
| out = list(boxes) | |
| for k, i in enumerate(order): | |
| x0, y0, x1, y1 = boxes[i] | |
| want = (y1 - y0) * PAD_FRAC | |
| up = y0 - boxes[order[k - 1]][3] if k > 0 else want * 2 | |
| dn = boxes[order[k + 1]][1] - y1 if k + 1 < len(order) else want * 2 | |
| out[i] = (x0, max(0.0, y0 - min(want, max(0.0, up) * 0.45)), | |
| x1, min(float(height), y1 + min(want * 1.4, max(0.0, dn) * 0.45))) | |
| return out | |
| def _forward(imgs): | |
| """One forward pass on a list of HxW grayscale crops -> list of T x C arrays. | |
| Each crop is decoded from its own T columns only, so sub-batches may pad to | |
| different widths and still be merged as a plain list. | |
| """ | |
| W = max(im.width for im in imgs) | |
| W = int(np.ceil(W / DOWNSAMPLE) * DOWNSAMPLE) | |
| x = torch.ones(len(imgs), 1, HEIGHT, W) | |
| for k, im in enumerate(imgs): | |
| a = (np.array(im, dtype=np.float32) / 255.0 - 0.5) / 0.5 | |
| x[k, 0, :, :im.width] = torch.from_numpy(a) | |
| with torch.no_grad(): | |
| ids = _model(x.to("cuda")).cpu().argmax(-1).numpy() | |
| return [ids[k][:max(1, im.width // DOWNSAMPLE)] for k, im in enumerate(imgs)] | |
| def _batched(part): | |
| """Recognise a batch, halving it if the GPU refuses the allocation. | |
| A very wide batch can exhaust the worker's memory pool (the CUDA caching | |
| allocator raises an NVML internal assert rather than a clean OOM); retrying | |
| with smaller sub-batches keeps a big page from failing outright. | |
| """ | |
| try: | |
| return _forward(part) | |
| except RuntimeError: | |
| if len(part) == 1: | |
| raise | |
| mid = len(part) // 2 | |
| return _batched(part[:mid]) + _batched(part[mid:]) | |
| def _recognise(crops): | |
| """Batched recognition (width-sorted so each batch pads to one width).""" | |
| texts = [] | |
| order = sorted(range(len(crops)), key=lambda i: crops[i].width) | |
| for s in range(0, len(order), BATCH): | |
| part = order[s:s + BATCH] | |
| rows = _batched([crops[i] for i in part]) | |
| for k, i in enumerate(part): | |
| prev, s_ = -1, [] | |
| for c in rows[k]: | |
| if c != prev and c != 0: | |
| s_.append(_itos[c]) | |
| prev = c | |
| texts.append((i, "".join(s_))) | |
| return [t for _, t in sorted(texts)] | |
| def _detect_polys(image): | |
| """Run the detector via a unique temp file (no fixed path: handlers run concurrently).""" | |
| with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as fh: | |
| image.convert("RGB").save(fh.name) | |
| tmp = fh.name | |
| try: | |
| polys = [] | |
| for res in _detector().predict(tmp): | |
| d = res if isinstance(res, dict) else (getattr(res, "json", {}) or {}).get("res", {}) | |
| polys += [np.array(p).reshape(-1, 2) for p in d.get("dt_polys", [])] | |
| return polys | |
| finally: | |
| try: | |
| os.remove(tmp) | |
| except OSError: | |
| pass | |
| # ------------------------------------------------------------------------ inference | |
| def read_line(image: Image.Image, pad: int = 8) -> str: | |
| """Read a single cropped line of Azerbaijani text (Latin or Cyrillic script). | |
| Args: | |
| image: a crop containing exactly one line of text, dark text on a light background. | |
| pad: extra white margin, in pixels, wrapped around the crop before recognition. | |
| """ | |
| if image is None: | |
| raise gr.Error("Upload an image of one text line first.") | |
| img = image.convert("L") | |
| pad = max(0, int(pad)) | |
| if pad: | |
| canvas = Image.new("L", (img.width + 2 * pad, img.height + 2 * pad), 255) | |
| canvas.paste(img, (pad, pad)) | |
| img = canvas | |
| with torch.no_grad(): | |
| logits = _model(preprocess(img).to("cuda")) | |
| return ctc_greedy(logits.cpu(), _itos) | |
| def read_page(image: Image.Image, col: str = "single", | |
| join_hyphens: bool = False) -> str: | |
| """Read a whole page scan: detect the lines, then recognise each one. | |
| Reading order assumes a single text column. For a two-column page, crop the | |
| left half in "Left column" mode and the right half in "Right column" mode | |
| (the columns are then read in the right order). | |
| Args: | |
| image: a page scan, roughly upright, at least ~150 dpi. | |
| col: "single" for a one-column page, "left" or "right" to isolate one column of a two-column page. | |
| join_hyphens: reattach words broken across a line break. | |
| """ | |
| if image is None: | |
| raise gr.Error("Upload a page image first.") | |
| img = image.convert("RGB") | |
| if col == "left": | |
| img = img.crop((0, 0, img.width // 2, img.height)) | |
| elif col == "right": | |
| img = img.crop((img.width // 2, 0, img.width, img.height)) | |
| polys = _detect_polys(img) | |
| if not polys: | |
| return "(no text detected)" | |
| polys = _reading_order(polys) | |
| boxes = [(float(p[:, 0].min()), float(p[:, 1].min()), | |
| float(p[:, 0].max()), float(p[:, 1].max())) for p in polys] | |
| boxes = _pad(boxes, img.height) | |
| crops = [] | |
| for x0, y0, x1, y1 in boxes: | |
| if (y1 - y0) >= (x1 - x0): # vertical text, not a line | |
| continue | |
| c = img.convert("L").crop((int(x0), int(y0), int(x1), int(y1))) | |
| if c.width < 8 or c.height < 8: | |
| continue | |
| r = HEIGHT / c.height | |
| c = c.resize((max(8, int(c.width * r)), HEIGHT), Image.LANCZOS) | |
| if c.width > MAX_WIDTH: | |
| c = c.crop((0, 0, MAX_WIDTH, HEIGHT)) | |
| crops.append(c) | |
| lines = _recognise(crops) | |
| if not join_hyphens: | |
| return "\n".join(lines) | |
| out = [] | |
| for ln in lines: | |
| if out and out[-1].endswith("-"): | |
| out[-1] = out[-1][:-1] + ln.lstrip() | |
| else: | |
| out.append(ln) | |
| return "\n".join(out) | |
| # ------------------------------------------------------------------------------- UI | |
| CSS = """ | |
| #col-container { max-width: 1100px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(title="Azerbaijani OCR") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# Azerbaijani OCR — CRNN + CTC\n" | |
| "An 8.5M-parameter recogniser for Azerbaijani print, **Latin and Cyrillic** " | |
| "script. It reads one line at a time; page mode pairs it with a text " | |
| "detector to do whole scans.\n\n" | |
| "[model card](https://huggingface.co/LocalDoc/azerbaijani-ocr-crnn) · " | |
| "[benchmark](https://huggingface.co/datasets/LocalDoc/azerbaijani-ocr-benchmark) " | |
| "(CER 0.0242 / WER 0.0565)" | |
| ) | |
| with gr.Tabs(): | |
| with gr.Tab("Line"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| line_in = gr.Image(type="pil", label="One text line", | |
| sources=["upload", "clipboard"]) | |
| line_btn = gr.Button("Read line", variant="primary") | |
| with gr.Accordion("Options", open=False): | |
| line_pad = gr.Slider(0, 24, value=8, step=1, | |
| label="White margin (px)") | |
| with gr.Column(scale=1): | |
| line_out = gr.Textbox(label="Recognised text", lines=4, | |
| buttons=["copy"]) | |
| with gr.Tab("Page"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| page_in = gr.Image(type="pil", label="Page scan", | |
| sources=["upload", "clipboard"]) | |
| page_btn = gr.Button("Read page", variant="primary") | |
| with gr.Accordion("Options", open=False): | |
| page_col = gr.Radio(["single", "left", "right"], | |
| value="single", | |
| label="Column layout", | |
| info="Two-column pages: read the left " | |
| "half, then the right half.") | |
| page_hyph = gr.Checkbox(value=False, | |
| label="Rejoin hyphenated words") | |
| with gr.Column(scale=1): | |
| page_out = gr.Textbox(label="Page text", lines=16, | |
| buttons=["copy"]) | |
| line_btn.click(read_line, inputs=[line_in, line_pad], outputs=line_out, | |
| api_name="read_line") | |
| page_btn.click(read_page, inputs=[page_in, page_col, page_hyph], | |
| outputs=page_out, api_name="read_page") | |
| gr.Markdown("### Examples") | |
| gr.Markdown( | |
| "Real line crops and pages from " | |
| "[LocalDoc/azerbaijani-ocr-benchmark](https://huggingface.co/datasets/LocalDoc/azerbaijani-ocr-benchmark) " | |
| "(CC-BY-4.0), the authors' own evaluation set." | |
| ) | |
| gr.Examples( | |
| examples=[["assets/line_latin_0.png"], ["assets/line_cyrillic_0.png"]], | |
| inputs=[line_in], outputs=line_out, fn=read_line, | |
| cache_examples=True, cache_mode="lazy", | |
| label="Line mode", | |
| ) | |
| gr.Examples( | |
| examples=[["assets/page_latin.png", "single", False], | |
| ["assets/page_cyrillic.png", "single", False]], | |
| inputs=[page_in, page_col, page_hyph], outputs=page_out, fn=read_page, | |
| cache_examples=True, cache_mode="lazy", | |
| label="Page mode", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |