import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces # noqa: E402 (must precede torch / CUDA-touching imports) import dataclasses # noqa: E402 import tempfile # noqa: E402 import time # noqa: E402 import gradio as gr # noqa: E402 import torch # noqa: E402 from PIL import Image, ImageDraw, ImageOps # noqa: E402 from huggingface_hub import hf_hub_download # noqa: E402 from kraken import serialization # noqa: E402 from kraken.configs import ( # noqa: E402 RecognitionInferenceConfig, SegmentationInferenceConfig, ) from kraken.tasks import RecognitionTaskModel, SegmentationTaskModel # noqa: E402 MODEL_ID = "small-models-for-glam/kraken-ppocrv6-medium" MODEL_FILE = "medium.safetensors" MODEL_PATH = hf_hub_download(MODEL_ID, MODEL_FILE) # Line detection: kraken's bundled BLLA baseline segmenter (`kraken segment -bl`). SEG_MODEL = SegmentationTaskModel.load_model() # Text recognition: the PP-OCRv6 medium CTC line recogniser. REC_MODEL = RecognitionTaskModel.load_model(MODEL_PATH) def _patch_batch_device(model) -> None: """Keep the CTC batch tensors on the network's device. kraken builds the per-line sequence lengths with the legacy ``torch.LongTensor([...])`` constructor, which escapes Fabric's ``init_tensor`` device context and therefore stays on the CPU. On CUDA that makes ``PPOCRv6Recognizer.forward`` compare a CPU ``out_lens`` against a CUDA ``positions`` tensor and raise "Expected all tensors to be on the same device". Re-homing the batch right before the forward pass fixes it without touching kraken itself. """ net = model.net # the wrapped PPOCRv6Model original = net._rec_predict def _rec_predict(line, lens=None): device = next(net.parameters()).device line = line.to(device) if lens is not None: lens = lens.to(device) return original(line, lens) net._rec_predict = _rec_predict _patch_batch_device(REC_MODEL) TEXT_DIRECTIONS = [ ("Left-to-right (Latin, Greek, Cyrillic…)", "horizontal-lr"), ("Right-to-left (Arabic, Hebrew, Syriac…)", "horizontal-rl"), ] CSS = """ #col-container { max-width: 1200px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ def _load_image(path: str, max_side: int) -> Image.Image: im = Image.open(path) im = ImageOps.exif_transpose(im) im = im.convert("RGB") if max(im.size) > max_side: im.thumbnail((max_side, max_side), Image.LANCZOS) return im def _draw_lines(im: Image.Image, segmentation) -> Image.Image: base = im.convert("RGBA") layer = Image.new("RGBA", im.size, (0, 0, 0, 0)) draw = ImageDraw.Draw(layer) width = max(2, im.width // 500) for line in segmentation.lines: if getattr(line, "boundary", None): pts = [(float(x), float(y)) for x, y in line.boundary] if len(pts) > 2: draw.polygon(pts, fill=(255, 128, 0, 45), outline=(255, 128, 0, 190)) if getattr(line, "baseline", None): pts = [(float(x), float(y)) for x, y in line.baseline] if len(pts) > 1: draw.line(pts, fill=(0, 132, 255, 235), width=width, joint="curve") return Image.alpha_composite(base, layer).convert("RGB") @spaces.GPU(duration=45) def transcribe( image: str, max_side: int = 1600, text_direction: str = "horizontal-lr", batch_size: int = 8, ) -> tuple: """Transcribe a document image with kraken PP-OCRv6 (medium). Detects text lines with kraken's baseline segmenter, then runs the multilingual PP-OCRv6 CTC recogniser over every detected line. Args: image: Path to the document image (manuscript page, printed page, scan). max_side: The image is downscaled so its longest side is at most this many pixels. text_direction: Principal reading direction of the page. batch_size: Number of text lines recognised per forward pass. Returns: A tuple of (line-detection overlay image, plain-text transcription, ALTO XML file path, markdown run statistics). """ if image is None: raise gr.Error("Please provide a document image first.") started = time.perf_counter() im = _load_image(image, int(max_side)) accelerator = "cuda" if torch.cuda.is_available() else "cpu" seg_config = SegmentationInferenceConfig( accelerator=accelerator, device="auto", precision="32-true", text_direction=text_direction, ) rec_config = RecognitionInferenceConfig( accelerator=accelerator, device="auto", precision="32-true", batch_size=int(batch_size), # ZeroGPU workers are daemonic forks and cannot spawn child processes, # so line extraction has to stay in-process. num_line_workers=0, ) seg_started = time.perf_counter() bounds = SEG_MODEL.predict(im, seg_config) seg_elapsed = time.perf_counter() - seg_started overlay = _draw_lines(im, bounds) if bounds.lines: rec_started = time.perf_counter() records = list(REC_MODEL.predict(im, bounds, rec_config)) rec_elapsed = time.perf_counter() - rec_started else: records = [] rec_elapsed = 0.0 text = "\n".join(record.prediction for record in records) results = dataclasses.replace(bounds, lines=records, imagename="input") alto = serialization.serialize( results=results, image_size=im.size, writing_mode="horizontal-tb", scripts=None, template="alto", template_source="native", ) tmp = tempfile.NamedTemporaryFile(prefix="kraken-", suffix=".alto.xml", delete=False, mode="w", encoding="utf-8") tmp.write(alto) tmp.close() confidences = [c for record in records for c in record.confidences] mean_conf = sum(confidences) / len(confidences) if confidences else 0.0 total = time.perf_counter() - started if records: stats = ( f"**{len(records)} text lines** · **{len(text)} characters** · " f"mean character confidence **{mean_conf:.1%}** \n" f"{im.size[0]}×{im.size[1]} px · segmentation {seg_elapsed:.1f}s · " f"recognition {rec_elapsed:.1f}s · total {total:.1f}s" ) else: stats = ( "**No text lines were detected.** Try a higher *Max image side* value in " "the advanced settings, or a scan with more visible text." ) return overlay, text, tmp.name, stats with gr.Blocks(title="kraken PP-OCRv6 medium") as demo: with gr.Column(elem_id="col-container"): gr.Markdown( "# kraken PP-OCRv6 · medium\n" "Multilingual OCR / handwritten-text recognition for historical and contemporary documents — " "**44 languages across 10 scripts** from a 15.9M-parameter CTC line recogniser.\n\n" "Text lines are detected with kraken's baseline segmenter, then transcribed line by line.\n\n" f"[Model card](https://huggingface.co/{MODEL_ID}) · [kraken](https://kraken.re)" ) with gr.Row(): with gr.Column(): image_in = gr.Image( label="Document image", type="filepath", sources=["upload", "clipboard"], height=420, ) run_btn = gr.Button("Transcribe", variant="primary") with gr.Accordion("Advanced settings", open=False): max_side = gr.Slider( 800, 3000, value=1600, step=100, label="Max image side (px)", info="Larger keeps more detail on dense pages but is slower.", ) text_direction = gr.Dropdown( choices=TEXT_DIRECTIONS, value="horizontal-lr", label="Reading direction", info="Affects line reading order on multi-column pages.", ) batch_size = gr.Slider(1, 16, value=8, step=1, label="Recognition batch size") with gr.Column(): seg_out = gr.Image(label="Detected text lines", height=420) text_out = gr.Textbox( label="Transcription", lines=18, max_lines=40, buttons=["copy"], placeholder="The transcription appears here, one line per detected text line.", ) stats_out = gr.Markdown() alto_out = gr.File(label="ALTO XML (line + word coordinates)") gr.Examples( examples=[ ["examples/printed-french-histoire.jpg"], ["examples/handwritten-spanish-notarial.jpg"], ["examples/printed-arabic.webp"], ], inputs=[image_in], outputs=[seg_out, text_out, alto_out, stats_out], fn=transcribe, cache_examples=True, cache_mode="lazy", label="Examples (from the kraken test suite)", ) gr.on( triggers=[run_btn.click], fn=transcribe, inputs=[image_in, max_side, text_direction, batch_size], outputs=[seg_out, text_out, alto_out, stats_out], api_name="transcribe", ) if __name__ == "__main__": demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)