"""Gradio application for layout-aware OCR with Rukopys-OCR-4B.""" import ipaddress import json import math import os import re import socket import time from io import BytesIO from urllib.parse import urljoin, urlparse # The default PyTorch kernel cache is read-only in ZeroGPU workers. os.environ.setdefault("PYTORCH_KERNEL_CACHE_PATH", "/tmp/torch-kernels") import gradio as gr import requests import spaces import torch from pdf2image import ( convert_from_bytes, convert_from_path, pdfinfo_from_bytes, pdfinfo_from_path, ) from PIL import Image, ImageDraw, ImageFont from transformers import AutoModelForMultimodalLM, AutoProcessor MODEL_ID = "ebinan92/Rukopys-OCR-4B" ATTENTION_IMPLEMENTATION = "sdpa" PROMPT = ( "Detect every text region in this Ukrainian handwritten document and " "return a JSON array of regions. Each region has bbox (x1 y1 x2 y2 in " "0..1000 normalized image coordinates), type (handwritten | printed | " "formula | table | annotation | image | graph), and text (transcription; " "empty for image/graph; LaTeX for formula; pipe-separated for table)." ) COLORS = { "handwritten": "#2563eb", "printed": "#16a34a", "formula": "#9333ea", "table": "#ea580c", "annotation": "#dc2626", "image": "#0891b2", "graph": "#ca8a04", } PDF_DPI = 180 MAX_RENDER_SIDE = 2600 MAX_REMOTE_BYTES = 25 * 1024 * 1024 MAX_IMAGE_PIXELS = 40_000_000 MAX_REDIRECTS = 4 REMOTE_TIMEOUT = (5, 30) JSON_FENCE_RE = re.compile(r"```(?:json)?\s*(.*?)```", re.DOTALL | re.IGNORECASE) NUMBER_RE = re.compile(r"[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?") VISION_PATCH_SIZE = 16 VISION_MERGE_SIZE = 2 VISION_FACTOR = VISION_PATCH_SIZE * VISION_MERGE_SIZE MODEL_MIN_PIXELS = 256 * VISION_FACTOR**2 MODEL_MAX_PIXELS = 4096 * VISION_FACTOR**2 type Region = dict[str, object] type PdfSource = str | bytes type SliderUpdate = dict[str, object] type PdfLoadResult = tuple[ Image.Image | None, PdfSource | None, SliderUpdate, str, ] Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS # ZeroGPU emulates CUDA while the app starts and replaces it with a real GPU # inside @spaces.GPU calls. Keeping the model on CUDA here avoids a costly # CPU-to-GPU transfer on every request. processor = AutoProcessor.from_pretrained( MODEL_ID, min_pixels=MODEL_MIN_PIXELS, max_pixels=MODEL_MAX_PIXELS, ) model = AutoModelForMultimodalLM.from_pretrained( MODEL_ID, dtype=torch.bfloat16, attn_implementation=ATTENTION_IMPLEMENTATION, ).to("cuda") active_attention = "SDPA" model.eval() print(f"Attention implementation: {active_attention}", flush=True) def _extract_json(text: str) -> object: """Parse a JSON value from plain output or a fenced model response.""" cleaned = text.strip() fenced = JSON_FENCE_RE.search(cleaned) if fenced: cleaned = fenced.group(1).strip() try: return json.loads(cleaned) except json.JSONDecodeError: start = min( (i for i in (cleaned.find("["), cleaned.find("{")) if i >= 0), default=-1 ) if start < 0: raise decoder = json.JSONDecoder() value, _ = decoder.raw_decode(cleaned[start:]) return value def _parse_regions(text: str) -> list[Region]: value = _extract_json(text) if not isinstance(value, list): raise ValueError("The model response is not a JSON array.") if not all(isinstance(region, dict) for region in value): raise ValueError("Every item in the region array must be a JSON object.") regions: list[Region] = [] for region in value: normalized = dict(region) bbox = _coerce_bbox(normalized.get("bbox")) if bbox is not None: normalized["bbox"] = list(bbox) regions.append(normalized) return regions def _coerce_bbox(value: object) -> tuple[float, float, float, float] | None: """Accept array boxes and common model strings such as '(x1,y1),(x2,y2)'.""" coordinates: list[object] if isinstance(value, (list, tuple)) and len(value) == 4: coordinates = list(value) elif isinstance(value, str): coordinates = NUMBER_RE.findall(value) if len(coordinates) != 4: return None else: return None try: bbox = tuple(float(coordinate) for coordinate in coordinates) except (TypeError, ValueError): return None if len(bbox) != 4 or not all(math.isfinite(coordinate) for coordinate in bbox): return None return bbox def _annotate(image: Image.Image, regions: list[Region]) -> Image.Image: result = image.convert("RGB").copy() draw = ImageDraw.Draw(result) font = ImageFont.load_default() width, height = result.size line_width = max(2, round(min(width, height) / 350)) for index, region in enumerate(regions, start=1): bbox = _coerce_bbox(region.get("bbox")) if bbox is None: continue x1, y1, x2, y2 = bbox left, right = sorted((x1 * width / 1000, x2 * width / 1000)) top, bottom = sorted((y1 * height / 1000, y2 * height / 1000)) coords = ( max(0, min(width, left)), max(0, min(height, top)), max(0, min(width, right)), max(0, min(height, bottom)), ) region_type = str(region.get("type", "region")) color = COLORS.get(region_type, "#475569") draw.rectangle(coords, outline=color, width=line_width) label = f"{index} · {region_type}" label_box = draw.textbbox( (coords[0], coords[1]), label, font=font, stroke_width=1 ) label_height = label_box[3] - label_box[1] + 6 label_width = label_box[2] - label_box[0] + 8 label_y = max(0, coords[1] - label_height) draw.rectangle( ( coords[0], label_y, min(width, coords[0] + label_width), label_y + label_height, ), fill=color, ) draw.text( (coords[0] + 4, label_y + 3), label, fill="white", font=font, stroke_width=1, stroke_fill=color, ) return result def _gpu_duration(_image: Image.Image, max_new_tokens: int) -> int: # A tighter estimate receives better queue priority while still leaving # enough headroom for dense documents. return min(300, max(60, int(max_new_tokens / 32) + 45)) def _prepare_image(image: Image.Image) -> Image.Image: """Convert an image to RGB and bound its dimensions for model input.""" image = image.convert("RGB") if max(image.size) > MAX_RENDER_SIDE: image.thumbnail((MAX_RENDER_SIDE, MAX_RENDER_SIDE), Image.Resampling.LANCZOS) return image def _vision_token_count(inputs: dict[str, torch.Tensor]) -> int | None: """Return the number of merged image tokens when processor metadata exists.""" grid = inputs.get("image_grid_thw") if not isinstance(grid, torch.Tensor) or grid.numel() == 0: return None return int(grid.prod(dim=-1).sum().item()) // VISION_MERGE_SIZE**2 def _render_pdf_page(source: PdfSource | None, page_number: int) -> Image.Image: if source is None: raise gr.Error("Upload a PDF first.") try: render = convert_from_bytes if isinstance(source, bytes) else convert_from_path pages = render( source, dpi=PDF_DPI, first_page=int(page_number), last_page=int(page_number), fmt="png", thread_count=1, timeout=60, size=MAX_RENDER_SIDE, ) except Exception as exc: raise gr.Error(f"Could not render PDF page {page_number}: {exc}") from exc if not pages: raise gr.Error(f"PDF page {page_number} could not be rendered.") return _prepare_image(pages[0]) def _pdf_page_count(source: PdfSource) -> int: try: inspect = pdfinfo_from_bytes if isinstance(source, bytes) else pdfinfo_from_path page_count = int(inspect(source, timeout=30).get("Pages", 0)) except Exception as exc: raise gr.Error(f"Could not read this PDF: {exc}") from exc if page_count < 1: raise gr.Error("The PDF contains no pages.") return page_count def _pdf_outputs(source: PdfSource) -> PdfLoadResult: page_count = _pdf_page_count(source) image = _render_pdf_page(source, 1) suffix = "s" if page_count != 1 else "" status = f"Loaded **{page_count} page{suffix}**. Page 1 is ready for OCR." slider = gr.update(minimum=1, maximum=page_count, value=1) return image, source, slider, status def load_pdf(pdf_file: str | None) -> PdfLoadResult: """Inspect a PDF and render its first page into the OCR image input.""" if pdf_file is None: slider = gr.update(minimum=1, maximum=1, value=1) return None, None, slider, "" return _pdf_outputs(pdf_file) def clear_pdf_source() -> tuple[None, SliderUpdate, str]: """Reset stale PDF navigation when the user uploads an image directly.""" slider = gr.update(minimum=1, maximum=1, value=1) return None, slider, "" def _validate_remote_url(url: str) -> None: try: parsed = urlparse(url) port = parsed.port except ValueError as exc: raise gr.Error("The URL contains an invalid port.") from exc if parsed.scheme not in {"http", "https"} or not parsed.hostname: raise gr.Error("Enter a valid HTTP or HTTPS URL.") if parsed.username or parsed.password: raise gr.Error("URLs containing credentials are not supported.") try: addresses = {item[4][0] for item in socket.getaddrinfo(parsed.hostname, port)} except socket.gaierror as exc: raise gr.Error("The URL hostname could not be resolved.") from exc if any(not ipaddress.ip_address(address).is_global for address in addresses): raise gr.Error( "Private, local, and reserved network addresses are not allowed." ) def _download_url(url: str) -> tuple[bytes, str]: current_url = url.strip() headers = {"User-Agent": "Rukopys-OCR-Space/1.0"} try: for _ in range(MAX_REDIRECTS + 1): _validate_remote_url(current_url) with requests.get( current_url, allow_redirects=False, headers=headers, stream=True, timeout=REMOTE_TIMEOUT, ) as response: if response.is_redirect or response.is_permanent_redirect: location = response.headers.get("Location") if not location: raise gr.Error("The URL returned an invalid redirect.") current_url = urljoin(current_url, location) continue response.raise_for_status() try: declared_size = int(response.headers.get("Content-Length", 0)) except ValueError as exc: raise gr.Error("The server returned an invalid file size.") from exc if declared_size > MAX_REMOTE_BYTES: raise gr.Error("The remote file exceeds the 25 MB limit.") content = bytearray() for chunk in response.iter_content(chunk_size=64 * 1024): content.extend(chunk) if len(content) > MAX_REMOTE_BYTES: raise gr.Error("The remote file exceeds the 25 MB limit.") media_type = ( response.headers.get("Content-Type", "").split(";", 1)[0].lower() ) return bytes(content), media_type except requests.RequestException as exc: raise gr.Error(f"Could not download the URL: {exc}") from exc raise gr.Error(f"The URL redirected more than {MAX_REDIRECTS} times.") def load_url(url: str) -> PdfLoadResult: """Download and prepare a public image or PDF URL.""" if not url.strip(): raise gr.Error("Paste an image or PDF URL first.") content, media_type = _download_url(url) if content.lstrip().startswith(b"%PDF-") or media_type == "application/pdf": return _pdf_outputs(content) try: image = Image.open(BytesIO(content)) if image.width * image.height > MAX_IMAGE_PIXELS: raise gr.Error("The remote image exceeds the 40-megapixel limit.") image.load() except (OSError, ValueError, Image.DecompressionBombError) as exc: raise gr.Error("The URL does not point to a supported image or PDF.") from exc image = _prepare_image(image) slider = gr.update(minimum=1, maximum=1, value=1) return image, None, slider, "Loaded a remote image. It is ready for OCR." def select_pdf_page( source: PdfSource | None, page_number: int, ) -> tuple[Image.Image, str]: """Render the selected page from the active PDF source.""" image = _render_pdf_page(source, int(page_number)) return image, f"Page **{int(page_number)}** is ready for OCR." @spaces.GPU(duration=_gpu_duration) def run_ocr( image: Image.Image | None, max_new_tokens: int, ) -> tuple[Image.Image, object, str, str]: """Recognize document regions and return annotated and structured output.""" if image is None: raise gr.Error("Upload a document image first.") started_at = time.perf_counter() image = _prepare_image(image) messages = [ { "role": "user", "content": [{"type": "image"}, {"type": "text", "text": PROMPT}], } ] text = processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, enable_thinking=False, ) prepared_at = time.perf_counter() inputs = processor(text=[text], images=[image], return_tensors="pt").to( model.device ) torch.cuda.synchronize() preprocessed_at = time.perf_counter() with torch.inference_mode(): output_ids = model.generate( **inputs, max_new_tokens=int(max_new_tokens), do_sample=False, ) torch.cuda.synchronize() generated_at = time.perf_counter() new_tokens = output_ids[:, inputs["input_ids"].shape[1] :] raw_output = processor.batch_decode(new_tokens, skip_special_tokens=True)[0] try: regions = _parse_regions(raw_output) annotated = _annotate(image, regions) structured_output: object = regions except (json.JSONDecodeError, ValueError) as exc: gr.Warning(f"The response could not be parsed as a region array: {exc}") annotated = image structured_output = {"unparsed_output": raw_output} finished_at = time.perf_counter() input_tokens = int(inputs["input_ids"].shape[1]) generated_tokens = int(new_tokens.shape[1]) vision_tokens = _vision_token_count(inputs) vision_summary = str(vision_tokens) if vision_tokens is not None else "unknown" timing = ( "**GPU job timing** (queue wait excluded) \n" f"Attention: `{active_attention}` · " f"Prepare: `{prepared_at - started_at:.2f}s` · " f"Preprocess + transfer: `{preprocessed_at - prepared_at:.2f}s` · " f"Generate: `{generated_at - preprocessed_at:.2f}s` · " f"Decode + annotate: `{finished_at - generated_at:.2f}s` · " f"Total: `{finished_at - started_at:.2f}s` \n" f"Input tokens: `{input_tokens}` · Vision tokens: `{vision_summary}` · " f"Generated tokens: `{generated_tokens}`" ) print( "OCR timing: " f"prepare={prepared_at - started_at:.2f}s, " f"preprocess_transfer={preprocessed_at - prepared_at:.2f}s, " f"generate={generated_at - preprocessed_at:.2f}s, " f"postprocess={finished_at - generated_at:.2f}s, " f"total={finished_at - started_at:.2f}s, " f"input_tokens={input_tokens}, vision_tokens={vision_summary}, " f"generated_tokens={generated_tokens}", flush=True, ) return annotated, structured_output, raw_output, timing CSS = """ .gradio-container { max-width: 1280px !important; } .hero { text-align: center; margin: 1.5rem auto 1rem; } .hero h1 { font-size: 2.25rem; margin-bottom: .35rem; } .hero p { color: var(--body-text-color-subdued); font-size: 1.05rem; } .primary-btn { min-height: 48px; } """ with gr.Blocks(css=CSS, title="Rukopys OCR") as demo: gr.HTML("""

✍️ Rukopys OCR

Ukrainian handwritten document recognition with layout-aware structured output.

""") with gr.Row(equal_height=True): with gr.Column(scale=1): input_image = gr.Image( type="pil", label="Document image", sources=["upload", "clipboard"], image_mode="RGB", height=560, ) with gr.Accordion("Or upload a PDF", open=False): pdf_file = gr.File( label="PDF document", file_types=[".pdf"], type="filepath", ) pdf_page = gr.Slider( minimum=1, maximum=1, value=1, step=1, label="PDF page", info="Choose one page at a time, then run OCR.", ) pdf_status = gr.Markdown() with gr.Accordion("Or paste an image/PDF URL", open=False): source_url = gr.Textbox( label="Public file URL", placeholder="https://example.com/document.pdf", ) load_url_button = gr.Button("Load URL") max_tokens = gr.Slider( minimum=512, maximum=8192, value=1024, step=512, label="Maximum output tokens", info=( "Increase for pages with many regions; larger values use " "more ZeroGPU quota." ), ) run_button = gr.Button( "Recognize document", variant="primary", elem_classes="primary-btn" ) timing_output = gr.Markdown() with gr.Column(scale=1): annotated_image = gr.Image(label="Detected regions", height=560) with gr.Tabs(): with gr.Tab("Structured JSON"): json_output = gr.JSON(label="Regions") with gr.Tab("Raw model output"): raw_output = gr.Code(label="Raw output", language="json", lines=18) gr.Markdown( "Bounding boxes use normalized `0..1000` coordinates. Region types include " "`handwritten`, `printed`, `formula`, `table`, `annotation`, `image`, " "and `graph`. " f"Model: [{MODEL_ID}](https://huggingface.co/{MODEL_ID})" ) pdf_source = gr.State() pdf_file.change( fn=load_pdf, inputs=pdf_file, outputs=[input_image, pdf_source, pdf_page, pdf_status], api_name=False, ) input_image.upload( fn=clear_pdf_source, inputs=None, outputs=[pdf_source, pdf_page, pdf_status], api_name=False, ) pdf_page.release( fn=select_pdf_page, inputs=[pdf_source, pdf_page], outputs=[input_image, pdf_status], api_name=False, ) load_url_button.click( fn=load_url, inputs=source_url, outputs=[input_image, pdf_source, pdf_page, pdf_status], api_name=False, ) source_url.submit( fn=load_url, inputs=source_url, outputs=[input_image, pdf_source, pdf_page, pdf_status], api_name=False, ) run_button.click( fn=run_ocr, inputs=[input_image, max_tokens], outputs=[annotated_image, json_output, raw_output, timing_output], api_name="recognize", ) demo.queue(default_concurrency_limit=1, max_size=20) if __name__ == "__main__": demo.launch()