lorensation/doc-benchmark-deepseek-ocr
0
1"""Streamlit UI for benchmarking DeepSeek-OCR vs Tesseract.2 3Features:4- Upload a single image with optional ground truth to run a benchmark.5- Run a quick batch benchmark on the bundled SROIE2019 dataset.6- Browse summaries of previous benchmark runs.7"""8 9from __future__ import annotations10 11import json12import os13import sys14import tempfile15import uuid16from datetime import datetime17from pathlib import Path18from typing import Dict, List, Optional, Tuple19 20import streamlit as st21 22# Ensure we can import the worker package when running in HF Spaces23REPO_ROOT = Path(__file__).resolve().parents[1]24if str(REPO_ROOT) not in sys.path:25 sys.path.append(str(REPO_ROOT))26 27from services.benchmark_worker import worker # noqa: E40228 29 30# ---------------------------------------------------------------------------31# Paths and configuration helpers32# ---------------------------------------------------------------------------33DATA_DIR = worker.DATA_DIR34UPLOAD_DIR = worker.UPLOAD_DIR35RESULTS_DIR = worker.RESULTS_DIR36SROIE_PATH = worker.SROIE_DATASET_PATH37 38 39def _ensure_dirs() -> None:40 UPLOAD_DIR.mkdir(parents=True, exist_ok=True)41 RESULTS_DIR.mkdir(parents=True, exist_ok=True)42 43 44def _save_upload(file) -> Path:45 """Persist an uploaded file to the uploads directory and return its path."""46 suffix = Path(file.name).suffix or ".png"47 dest = UPLOAD_DIR / f"{uuid.uuid4()}{suffix}"48 with open(dest, "wb") as fh:49 fh.write(file.getbuffer())50 return dest51 52 53def _parse_ground_truth(gt_text: str, gt_json_file) -> Tuple[Optional[str], Optional[List[str]], Optional[Dict]]:54 """Normalize ground truth inputs from the UI."""55 ground_truth_text = gt_text.strip() if gt_text and gt_text.strip() else None56 ground_truth_lines = ground_truth_text.splitlines() if ground_truth_text else None57 58 ground_truth_fields = None59 if gt_json_file is not None:60 try:61 ground_truth_fields = json.loads(gt_json_file.getvalue().decode("utf-8"))62 except Exception as exc: # pragma: no cover - UI feedback path63 st.error(f"Could not parse ground truth JSON: {exc}")64 65 return ground_truth_text, ground_truth_lines, ground_truth_fields66 67 68# ---------------------------------------------------------------------------69# Rendering helpers70# ---------------------------------------------------------------------------71def _render_result_details(result: dict) -> None:72 """Pretty-print a single benchmark result."""73 st.success("Benchmark complete")74 meta_cols = st.columns(3)75 meta_cols[0].metric("Run ID", result.get("run_id", "") or "-")76 meta_cols[1].metric("File", result.get("filename", "") or "-")77 meta_cols[2].metric("Created", result.get("created_at", "") or "-")78 79 st.caption(f"Saved to: {result.get('result_file', 'not saved')}")80 81 metrics = result.get("metrics", {})82 lengths = metrics.get("lengths") or {83 "deepseek": len(result.get("deepseek_text", "")),84 "tesseract": len(result.get("tesseract_text", "")),85 }86 char_diff = metrics.get("char_difference") or abs(lengths["deepseek"] - lengths["tesseract"])87 88 with st.container():89 cols = st.columns(3)90 cols[0].metric("DeepSeek length", lengths.get("deepseek", 0))91 cols[1].metric("Tesseract length", lengths.get("tesseract", 0))92 cols[2].metric("Char diff", char_diff)93 94 st.divider()95 96 deepseek_metrics = metrics.get("deepseek", {})97 tesseract_metrics = metrics.get("tesseract", {})98 99 if deepseek_metrics or tesseract_metrics:100 st.subheader("Quality metrics")101 mcols = st.columns(2)102 _render_model_metrics("DeepSeek-OCR", deepseek_metrics, mcols[0])103 _render_model_metrics("Tesseract", tesseract_metrics, mcols[1])104 105 st.subheader("Extracted text")106 text_cols = st.columns(2)107 text_cols[0].write("DeepSeek-OCR")108 text_cols[0].code(result.get("deepseek_text", ""), language="markdown")109 110 text_cols[1].write("Tesseract")111 text_cols[1].code(result.get("tesseract_text", ""), language="markdown")112 113 st.download_button(114 "Download result JSON",115 data=json.dumps(result, ensure_ascii=False, indent=2),116 file_name=f"{result.get('run_id', 'benchmark')}.json",117 mime="application/json",118 use_container_width=True,119 )120 121 122def _render_model_metrics(name: str, payload: dict, container) -> None:123 """Render key metrics for a single model."""124 if not payload:125 container.info("No metrics captured.")126 return127 128 text_metrics = payload.get("text", {})129 llm_metrics = payload.get("llm", {})130 131 if text_metrics:132 cer = text_metrics.get("cer", {}).get("cer")133 wer = text_metrics.get("wer", {}).get("wer")134 ser = text_metrics.get("ser", {}).get("ser")135 if cer is not None or wer is not None or ser is not None:136 container.metric(f"{name} CER", f"{cer:.2f}%" if cer is not None else "-")137 container.metric(f"{name} WER", f"{wer:.2f}%" if wer is not None else "-")138 container.metric(f"{name} SER", f"{ser:.2f}%" if ser is not None else "-")139 140 if llm_metrics:141 token_eff = llm_metrics.get("token_efficiency", {})142 fext = llm_metrics.get("field_extraction", {})143 rows = []144 if token_eff:145 rows.append(f"Token density: {token_eff.get('token_density', 0):.2f}")146 rows.append(f"Chars/token: {token_eff.get('chars_per_token', 0):.2f}")147 if fext:148 rows.append(f"Field F1: {fext.get('field_f1', 0):.2f}")149 rows.append(f"Field recall: {fext.get('field_recall', 0):.2f}")150 if rows:151 container.write("\n".join(rows))152 153 154def _render_batch_summary(summary: dict, results: List[dict]) -> None:155 st.success("SROIE benchmark finished")156 st.write(157 f"Batch ID: `{summary.get('batch_run_id', '-')}` | "158 f"Split: `{summary.get('split', '-')}` | "159 f"Samples: `{summary.get('count', 0)}`"160 )161 if summary.get("result_files"):162 st.caption(f"Stored in: {Path(summary['result_files'][0]).parent}")163 164 if results:165 st.subheader("Individual results")166 for item in results:167 with st.expander(f"{item.get('filename', 'sample')} — {item.get('run_id', '')}"):168 _render_result_details(item)169 else:170 st.info("No individual results were returned.")171 172 173# ---------------------------------------------------------------------------174# Page renderers175# ---------------------------------------------------------------------------176def page_single_run() -> None:177 st.header("Upload and Benchmark")178 st.write(179 "Upload an image and optional ground truth. The worker will call DeepSeek-OCR "180 "and Tesseract services, compute metrics, and save the result to the workspace."181 )182 183 upload = st.file_uploader(184 "Image file",185 type=["png", "jpg", "jpeg", "tiff", "bmp", "webp"],186 )187 gt_text = st.text_area(188 "Ground truth text (optional)",189 placeholder="Paste expected text. Leave empty to auto-use SROIE ground truth when the file belongs to the dataset.",190 height=120,191 )192 gt_json = st.file_uploader(193 "Ground truth fields JSON (optional)",194 type=["json"],195 key="gt_json_single",196 help='Optional key-value ground truth (e.g. {"total": "123.00", "date": "2024-01-01"}).',197 )198 199 if upload and st.button("Run benchmark", type="primary", use_container_width=True):200 _ensure_dirs()201 with st.spinner("Running benchmark..."):202 img_path = _save_upload(upload)203 gt_text_val, gt_lines, gt_fields = _parse_ground_truth(gt_text, gt_json)204 result = worker.run_benchmark(205 str(img_path),206 ground_truth_text=gt_text_val,207 ground_truth_lines=gt_lines,208 ground_truth_fields=gt_fields,209 )210 _render_result_details(result)211 elif not upload:212 st.info("Upload an image to begin.")213 214 215def _count_sroie_samples(split: str) -> int:216 patterns = ("*.jpg", "*.jpeg", "*.png")217 img_dir = Path(SROIE_PATH) / split / "img"218 total = 0219 for pattern in patterns:220 total += len(list(img_dir.glob(pattern)))221 return total222 223 224def page_sroie_batch() -> None:225 st.header("SROIE2019 Batch Benchmark")226 st.write(227 "Run a quick benchmark on the bundled SROIE2019 dataset. "228 "Ground truth is auto-loaded from the dataset annotations."229 )230 231 if not SROIE_PATH.exists():232 st.error(f"SROIE dataset not found at {SROIE_PATH}. Add the dataset before running.")233 return234 235 split = st.selectbox("Split", ["train", "test"], index=0)236 total = _count_sroie_samples(split)237 st.caption(f"Found {total} images in {SROIE_PATH / split / 'img'}")238 limit = st.slider("Number of samples to run", min_value=1, max_value=max(1, total) if total else 1, value=min(3, total or 1))239 240 if st.button("Run SROIE benchmark", type="primary", use_container_width=True):241 _ensure_dirs()242 with st.spinner("Processing SROIE samples..."):243 payload = worker.run_sroie_samples(split=split, limit=limit)244 summary = payload.get("summary", {})245 results = payload.get("results", [])246 _render_batch_summary(summary, results)247 248 249def _is_batch_summary(data: dict) -> bool:250 return "batch_run_id" in data or ("summary" in data and isinstance(data["summary"], dict))251 252 253def _load_result_file(path: Path) -> Optional[dict]:254 try:255 with path.open("r", encoding="utf-8") as fh:256 return json.load(fh)257 except Exception:258 return None259 260 261def _collect_history(max_files: int = 200) -> Tuple[List[dict], List[dict]]:262 run_items: List[dict] = []263 batch_items: List[dict] = []264 265 files = sorted(RESULTS_DIR.rglob("*.json"), key=os.path.getmtime, reverse=True)266 for path in files[:max_files]:267 data = _load_result_file(path)268 if not data:269 continue270 271 if _is_batch_summary(data):272 summary = data.get("summary", data)273 batch_items.append(274 {275 "path": str(path),276 "batch_run_id": summary.get("batch_run_id"),277 "split": summary.get("split"),278 "count": summary.get("count"),279 "saved": summary.get("result_files", []),280 }281 )282 elif "run_id" in data:283 run_items.append(284 {285 "path": str(path),286 "run_id": data.get("run_id"),287 "filename": data.get("filename"),288 "created_at": data.get("created_at"),289 "ground_truth": data.get("ground_truth_available", bool(data.get("metrics", {}).get("deepseek", {}).get("text"))),290 "char_diff": data.get("metrics", {}).get("char_difference"),291 }292 )293 294 return run_items, batch_items295 296 297def page_history() -> None:298 st.header("Benchmark History")299 st.write("Browse results saved under the workspace `data/results` directory.")300 301 run_items, batch_items = _collect_history()302 303 if batch_items:304 st.subheader("Batch summaries")305 st.table(batch_items)306 else:307 st.info("No batch summaries found.")308 309 st.subheader("Individual runs")310 if not run_items:311 st.info("No individual runs found.")312 return313 314 st.dataframe(run_items, use_container_width=True)315 316 selected = st.selectbox(317 "Inspect a run",318 options=run_items,319 format_func=lambda item: f"{item['run_id']} — {item['filename']}",320 )321 322 if selected:323 data = _load_result_file(Path(selected["path"]))324 if data:325 _render_result_details(data)326 else:327 st.error("Could not load the selected result file.")328 329 330# ---------------------------------------------------------------------------331# App entrypoint332# ---------------------------------------------------------------------------333def main() -> None:334 st.set_page_config(page_title="DeepSeek OCR Benchmark", layout="wide")335 st.title("DeepSeek-OCR Benchmark Space")336 st.caption("Compare DeepSeek-OCR against Tesseract, run SROIE batches, and review saved results.")337 338 page = st.sidebar.radio(339 "Navigation",340 options=[341 "Upload benchmark",342 "SROIE dataset benchmark",343 "Benchmark history",344 ],345 )346 st.sidebar.markdown(f"**Data dir:** `{DATA_DIR}`")347 st.sidebar.markdown(f"**Results dir:** `{RESULTS_DIR}`")348 st.sidebar.markdown(f"**Uploads dir:** `{UPLOAD_DIR}`")349 st.sidebar.markdown(f"**SROIE path:** `{SROIE_PATH}`")350 351 if page == "Upload benchmark":352 page_single_run()353 elif page == "SROIE dataset benchmark":354 page_sroie_batch()355 else:356 page_history()357 358 359if __name__ == "__main__":360 main()361