Spaces:
Build error
Build error
| """ | |
| Surya OCR 2 — Gradio UI | |
| CPU inference via llama.cpp backend. | |
| """ | |
| import os | |
| import json | |
| import re | |
| import tempfile | |
| from pathlib import Path | |
| # Must be set before importing surya | |
| os.environ.setdefault("SURYA_INFERENCE_BACKEND", "llamacpp") | |
| os.environ.setdefault("HF_HOME", "/tmp/hf_cache") | |
| import gradio as gr | |
| from PIL import Image, ImageDraw, ImageFont | |
| # ── Label → colour map ──────────────────────────────────────────────────────── | |
| LABEL_COLORS: dict[str, str] = { | |
| "Text": "#2196F3", | |
| "SectionHeader": "#9C27B0", | |
| "Table": "#FF9800", | |
| "Equation": "#F44336", | |
| "Picture": "#4CAF50", | |
| "Figure": "#4CAF50", | |
| "Form": "#00BCD4", | |
| "PageHeader": "#607D8B", | |
| "PageFooter": "#607D8B", | |
| "ListGroup": "#8BC34A", | |
| "Caption": "#795548", | |
| "Footnote": "#9E9E9E", | |
| "Code": "#FF5722", | |
| "TableOfContents": "#3F51B5", | |
| "Bibliography": "#795548", | |
| } | |
| # ── Lazy-loaded inference manager (spawns llama-server once) ───────────────── | |
| _manager = None | |
| def get_manager(): | |
| global _manager | |
| if _manager is None: | |
| from surya.inference import SuryaInferenceManager # noqa: PLC0415 | |
| _manager = SuryaInferenceManager() | |
| return _manager | |
| # ── PDF / image helpers ─────────────────────────────────────────────────────── | |
| def pdf_to_images(pdf_path: str, dpi: int = 150) -> list[Image.Image]: | |
| import fitz # PyMuPDF | |
| doc = fitz.open(pdf_path) | |
| scale = dpi / 72.0 | |
| mat = fitz.Matrix(scale, scale) | |
| images = [] | |
| for page in doc: | |
| pix = page.get_pixmap(matrix=mat, alpha=False) | |
| img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) | |
| images.append(img) | |
| doc.close() | |
| return images | |
| def parse_page_range(spec: str, total: int) -> list[int]: | |
| """Convert a page-range string like '0,2-5,8' into a sorted list of indices.""" | |
| if not spec.strip(): | |
| return list(range(total)) | |
| indices: set[int] = set() | |
| for part in spec.split(","): | |
| part = part.strip() | |
| m = re.match(r"^(\d+)-(\d+)$", part) | |
| if m: | |
| indices.update(range(int(m.group(1)), int(m.group(2)) + 1)) | |
| elif part.isdigit(): | |
| indices.add(int(part)) | |
| return sorted(i for i in indices if 0 <= i < total) | |
| def load_images(file_path: str, page_range_str: str) -> list[Image.Image]: | |
| if not file_path: | |
| return [] | |
| suffix = Path(file_path).suffix.lower() | |
| if suffix == ".pdf": | |
| all_pages = pdf_to_images(file_path) | |
| indices = parse_page_range(page_range_str, len(all_pages)) | |
| return [all_pages[i] for i in indices] | |
| else: | |
| return [Image.open(file_path).convert("RGB")] | |
| # ── Drawing helpers ─────────────────────────────────────────────────────────── | |
| def draw_boxes(img: Image.Image, boxes: list[dict], show_labels: bool = True) -> Image.Image: | |
| out = img.copy() | |
| draw = ImageDraw.Draw(out, "RGBA") | |
| try: | |
| font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 14) | |
| except Exception: | |
| font = ImageFont.load_default() | |
| for box in boxes: | |
| bbox = box.get("bbox") | |
| label = box.get("label", "?") | |
| order = box.get("reading_order") | |
| if not bbox: | |
| continue | |
| color = LABEL_COLORS.get(label, "#888888") | |
| r, g, b = int(color[1:3], 16), int(color[3:5], 16), int(color[5:7], 16) | |
| draw.rectangle(bbox, outline=(r, g, b, 255), width=2) | |
| draw.rectangle(bbox, fill=(r, g, b, 25)) | |
| if show_labels: | |
| tag = f"{order}:{label}" if order is not None else label | |
| tx, ty = bbox[0] + 2, max(0, bbox[1] - 16) | |
| draw.rectangle([tx - 1, ty - 1, tx + len(tag) * 8, ty + 16], fill=(r, g, b, 200)) | |
| draw.text((tx, ty), tag, fill="white", font=font) | |
| return out | |
| # ── Strip HTML to plain text ────────────────────────────────────────────────── | |
| def html_to_text(html: str) -> str: | |
| text = re.sub(r"<[^>]+>", " ", html or "") | |
| return re.sub(r"\s+", " ", text).strip() | |
| # ── OCR ─────────────────────────────────────────────────────────────────────── | |
| def run_ocr(file_obj, page_range: str, gr_progress=gr.Progress()): | |
| if file_obj is None: | |
| return "Upload a file first.", "", "{}", [] | |
| file_path = file_obj.name if hasattr(file_obj, "name") else file_obj | |
| gr_progress(0.05, desc="Loading pages…") | |
| images = load_images(file_path, page_range) | |
| if not images: | |
| return "No pages found.", "", "{}", [] | |
| gr_progress(0.15, desc=f"Starting inference on {len(images)} page(s)…") | |
| manager = get_manager() | |
| from surya.recognition import RecognitionPredictor # noqa: PLC0415 | |
| predictor = RecognitionPredictor(manager) | |
| text_pages, html_pages, json_pages, annotated = [], [], [], [] | |
| for idx, img in enumerate(images): | |
| gr_progress( | |
| 0.15 + 0.80 * (idx / len(images)), | |
| desc=f"OCR: page {idx + 1}/{len(images)}…", | |
| ) | |
| preds = predictor([img]) | |
| pred = preds[0] | |
| blocks_raw = [] | |
| text_chunks, html_chunks = [], [] | |
| for blk in getattr(pred, "blocks", []): | |
| if getattr(blk, "skipped", False): | |
| continue | |
| blk_html = getattr(blk, "html", "") or "" | |
| blk_text = html_to_text(blk_html) | |
| blk_label = getattr(blk, "label", "Text") | |
| blk_bbox = getattr(blk, "bbox", None) | |
| blk_conf = getattr(blk, "confidence", 1.0) | |
| blk_order = getattr(blk, "reading_order", None) | |
| if blk_text: | |
| text_chunks.append(blk_text) | |
| if blk_html: | |
| html_chunks.append(blk_html) | |
| blocks_raw.append({ | |
| "label": blk_label, | |
| "reading_order": blk_order, | |
| "text": blk_text, | |
| "html": blk_html, | |
| "bbox": blk_bbox, | |
| "confidence": round(float(blk_conf), 3), | |
| }) | |
| text_pages.append(f"═══ Page {idx + 1} ═══\n" + "\n\n".join(text_chunks)) | |
| html_pages.append(f"<!-- Page {idx + 1} -->\n" + "\n".join(html_chunks)) | |
| json_pages.append({"page": idx + 1, "blocks": blocks_raw}) | |
| annotated.append(draw_boxes(img, blocks_raw)) | |
| gr_progress(1.0, desc="Done!") | |
| return ( | |
| "\n\n".join(text_pages), | |
| "\n\n".join(html_pages), | |
| json.dumps(json_pages, indent=2, ensure_ascii=False), | |
| annotated, | |
| ) | |
| # ── Layout analysis ─────────────────────────────────────────────────────────── | |
| def run_layout(file_obj, page_range: str, gr_progress=gr.Progress()): | |
| if file_obj is None: | |
| return "Upload a file first.", "{}", [] | |
| file_path = file_obj.name if hasattr(file_obj, "name") else file_obj | |
| gr_progress(0.05, desc="Loading pages…") | |
| images = load_images(file_path, page_range) | |
| if not images: | |
| return "No pages found.", "{}", [] | |
| gr_progress(0.15, desc="Running layout analysis…") | |
| manager = get_manager() | |
| from surya.layout import LayoutPredictor # noqa: PLC0415 | |
| predictor = LayoutPredictor(manager) | |
| predictions = predictor(images) | |
| summary_parts, json_pages, annotated = [], [], [] | |
| for idx, (img, pred) in enumerate(zip(images, predictions)): | |
| bboxes = getattr(pred, "bboxes", []) | |
| counts: dict[str, int] = {} | |
| boxes_raw = [] | |
| for box in bboxes: | |
| label = getattr(box, "label", "Unknown") | |
| bbox = getattr(box, "bbox", None) | |
| position = getattr(box, "position", None) | |
| confidence = getattr(box, "confidence", 1.0) | |
| counts[label] = counts.get(label, 0) + 1 | |
| boxes_raw.append({ | |
| "label": label, | |
| "reading_order": position, | |
| "bbox": bbox, | |
| "confidence": round(float(confidence), 3), | |
| }) | |
| lines = [f" {k}: {v}" for k, v in sorted(counts.items())] | |
| summary_parts.append(f"═══ Page {idx + 1} ═══\n" + "\n".join(lines)) | |
| json_pages.append({"page": idx + 1, "layout": boxes_raw}) | |
| annotated.append(draw_boxes(img, boxes_raw)) | |
| gr_progress(1.0, desc="Done!") | |
| return ( | |
| "\n\n".join(summary_parts), | |
| json.dumps(json_pages, indent=2), | |
| annotated, | |
| ) | |
| # ── Table recognition ───────────────────────────────────────────────────────── | |
| def run_table_rec(file_obj, page_range: str, gr_progress=gr.Progress()): | |
| if file_obj is None: | |
| return "Upload a file first.", "{}" | |
| file_path = file_obj.name if hasattr(file_obj, "name") else file_obj | |
| gr_progress(0.05, desc="Loading pages…") | |
| images = load_images(file_path, page_range) | |
| if not images: | |
| return "No pages found.", "{}" | |
| gr_progress(0.15, desc="Running table recognition…") | |
| manager = get_manager() | |
| from surya.table_rec import TableRecPredictor # noqa: PLC0415 | |
| predictor = TableRecPredictor(manager) | |
| predictions = predictor(images) | |
| html_parts, json_pages = [], [] | |
| for idx, pred in enumerate(predictions): | |
| tbl_html = getattr(pred, "html", "") or "" | |
| rows = len(getattr(pred, "rows", [])) | |
| cols = len(getattr(pred, "cols", [])) | |
| html_parts.append( | |
| f"<h3 style='margin:8px 0'>Page {idx + 1}</h3>" | |
| + (tbl_html if tbl_html else "<p><em>No table detected.</em></p>") | |
| ) | |
| json_pages.append({"page": idx + 1, "rows": rows, "cols": cols, "html": tbl_html}) | |
| gr_progress(1.0, desc="Done!") | |
| return ( | |
| "<div style='font-family:sans-serif'>" + "\n".join(html_parts) + "</div>", | |
| json.dumps(json_pages, indent=2), | |
| ) | |
| # ── JSON download helper ────────────────────────────────────────────────────── | |
| def save_json(json_str: str): | |
| if not json_str or json_str in ("{}", ""): | |
| return None | |
| tmp = tempfile.NamedTemporaryFile(suffix=".json", delete=False, mode="w") | |
| tmp.write(json_str) | |
| tmp.close() | |
| return tmp.name | |
| def save_text(text: str): | |
| if not text: | |
| return None | |
| tmp = tempfile.NamedTemporaryFile(suffix=".txt", delete=False, mode="w") | |
| tmp.write(text) | |
| tmp.close() | |
| return tmp.name | |
| # ── Gradio UI ───────────────────────────────────────────────────────────────── | |
| CSS = """ | |
| #header { text-align: center; margin-bottom: 4px; } | |
| #header h1 { font-size: 2.2rem; font-weight: 700; } | |
| #header p { color: #6b7280; font-size: 1rem; } | |
| .warn-box { background:#fef3c7; border:1px solid #f59e0b; border-radius:8px; | |
| padding:10px 14px; color:#92400e; font-size:0.9rem; } | |
| """ | |
| INTRO_MD = """ | |
| <div id="header"> | |
| <h1>📄 Surya OCR 2</h1> | |
| <p>State-of-the-art document intelligence · 650M params · CPU inference</p> | |
| </div> | |
| """ | |
| WARN_MD = """ | |
| <div class="warn-box"> | |
| ⏳ <strong>CPU mode:</strong> The model downloads on first run (~400 MB GGUF). | |
| Inference is ~0.1 pages/second — a 5-page PDF takes ~50 s. | |
| Subsequent runs are faster once the model is cached. | |
| </div> | |
| """ | |
| with gr.Blocks(title="Surya OCR 2", css=CSS, theme=gr.themes.Soft()) as demo: | |
| gr.HTML(INTRO_MD) | |
| gr.HTML(WARN_MD) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| file_input = gr.File( | |
| label="Upload PDF or Image", | |
| file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp", ".webp"], | |
| ) | |
| page_range_input = gr.Textbox( | |
| label="Page range (PDF only, 0-indexed)", | |
| placeholder="e.g. 0,2-5 — leave blank for all pages", | |
| value="", | |
| ) | |
| with gr.Column(scale=2): | |
| with gr.Tabs(): | |
| # ── Tab 1: Full OCR ────────────────────────────────────────── | |
| with gr.Tab("📝 Full OCR"): | |
| ocr_btn = gr.Button("Run Full OCR", variant="primary", size="lg") | |
| with gr.Tabs(): | |
| with gr.Tab("Text"): | |
| ocr_text_out = gr.Textbox( | |
| label="Extracted text", | |
| lines=20, | |
| show_copy_button=True, | |
| ) | |
| ocr_text_dl = gr.File(label="Download .txt", interactive=False) | |
| with gr.Tab("HTML"): | |
| ocr_html_out = gr.HTML(label="HTML output") | |
| with gr.Tab("JSON"): | |
| ocr_json_out = gr.Code( | |
| label="Structured JSON", language="json", lines=20 | |
| ) | |
| ocr_json_dl = gr.File(label="Download .json", interactive=False) | |
| with gr.Tab("Annotated pages"): | |
| ocr_gallery = gr.Gallery( | |
| label="Pages with detected blocks", | |
| columns=2, | |
| height="auto", | |
| ) | |
| def _ocr(file, pr, prog=gr.Progress()): | |
| text, html, js, imgs = run_ocr(file, pr, prog) | |
| return text, html, js, imgs, save_text(text), save_json(js) | |
| ocr_btn.click( | |
| _ocr, | |
| inputs=[file_input, page_range_input], | |
| outputs=[ | |
| ocr_text_out, ocr_html_out, ocr_json_out, | |
| ocr_gallery, ocr_text_dl, ocr_json_dl, | |
| ], | |
| ) | |
| # ── Tab 2: Layout analysis ─────────────────────────────────── | |
| with gr.Tab("🗂 Layout Analysis"): | |
| lay_btn = gr.Button("Analyse Layout", variant="primary", size="lg") | |
| with gr.Tabs(): | |
| with gr.Tab("Summary"): | |
| lay_summary_out = gr.Textbox( | |
| label="Layout summary", | |
| lines=15, | |
| show_copy_button=True, | |
| ) | |
| with gr.Tab("JSON"): | |
| lay_json_out = gr.Code( | |
| label="Layout JSON", language="json", lines=20 | |
| ) | |
| lay_json_dl = gr.File(label="Download .json", interactive=False) | |
| with gr.Tab("Annotated pages"): | |
| lay_gallery = gr.Gallery( | |
| label="Pages with layout boxes", | |
| columns=2, | |
| height="auto", | |
| ) | |
| def _layout(file, pr, prog=gr.Progress()): | |
| summary, js, imgs = run_layout(file, pr, prog) | |
| return summary, js, imgs, save_json(js) | |
| lay_btn.click( | |
| _layout, | |
| inputs=[file_input, page_range_input], | |
| outputs=[lay_summary_out, lay_json_out, lay_gallery, lay_json_dl], | |
| ) | |
| # ── Tab 3: Table recognition ───────────────────────────────── | |
| with gr.Tab("📊 Table Recognition"): | |
| tbl_btn = gr.Button("Recognise Tables", variant="primary", size="lg") | |
| with gr.Tabs(): | |
| with gr.Tab("HTML table"): | |
| tbl_html_out = gr.HTML(label="Extracted tables") | |
| with gr.Tab("JSON"): | |
| tbl_json_out = gr.Code( | |
| label="Table JSON", language="json", lines=15 | |
| ) | |
| tbl_json_dl = gr.File(label="Download .json", interactive=False) | |
| def _table(file, pr, prog=gr.Progress()): | |
| html, js = run_table_rec(file, pr, prog) | |
| return html, js, save_json(js) | |
| tbl_btn.click( | |
| _table, | |
| inputs=[file_input, page_range_input], | |
| outputs=[tbl_html_out, tbl_json_out, tbl_json_dl], | |
| ) | |
| gr.Markdown( | |
| """ | |
| --- | |
| **Powered by [Surya OCR 2](https://huggingface.co/datalab-to/surya-ocr-2)** | |
| · 650 M params · Apache 2.0 code · OpenRail weights | |
| · Inference: llama.cpp (CPU) | |
| """, | |
| elem_id="footer", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=int(os.environ.get("GRADIO_SERVER_PORT", 7860)), | |
| show_error=True, | |
| ) | |