"""
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"\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"Page {idx + 1}
"
+ (tbl_html if tbl_html else "No table detected.
")
)
json_pages.append({"page": idx + 1, "rows": rows, "cols": cols, "html": tbl_html})
gr_progress(1.0, desc="Done!")
return (
"" + "\n".join(html_parts) + "
",
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 = """
"""
WARN_MD = """
⏳ CPU mode: 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.
"""
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,
)