Speed up Smoke Signal OCR and add live progress

smoke_signal_tab.py

@@ -708,37 +708,23 @@ def _profile_status_html() -> str:
 
 
 # ── Step 3: OCR ────────────────────────────────────────────────────────────────
-def run_ocr() -> tuple:
-    """Run Surya OCR on all profiled PDFs."""
-    df = load_manifest_df()
-    debug = f"[DEBUG] SS_ROOT={SS_ROOT}\nMANIFEST_CSV={MANIFEST_CSV}\nCSV exists={MANIFEST_CSV.exists()}\n"
-    if not df.empty:
-        debug += f"Manifest rows={len(df)}\nStatuses={df['status'].value_counts().to_dict()}\n"
-    else:
-        debug += "Manifest is EMPTY\n"
+_SURYA_RUNTIME = None
+try:
+    SS_SURYA_BATCH_SIZE = max(1, int(os.environ.get("SS_SURYA_BATCH_SIZE", "4")))
+except Exception:
+    SS_SURYA_BATCH_SIZE = 4
 
-    if df.empty:
-        return _ocr_status_html(), debug + "No sources. Complete Steps 1-2 first."
-
-    eligible = df[df["status"].isin(["profiled", "ocred", "rendered"])]
-    if eligible.empty:
-        unknown_or_excluded = df[df["rights_class"].isin(["unknown", "excluded"])] if "rights_class" in df.columns else df.iloc[0:0]
-        pending = df[df["status"] == "pending"] if "status" in df.columns else df.iloc[0:0]
-        debug += (
-            f"Eligible rows={len(eligible)}\n"
-            f"Pending rows={len(pending)}\n"
-            f"Unknown/excluded rights={len(unknown_or_excluded)}\n"
-            "Tip: In Step 1, update rights_class away from unknown/excluded, "
-            "then rerun Step 2 Profile.\n"
-        )
-        return _ocr_status_html(), debug + "No profiled PDFs. Complete Step 2 first."
 
-    log = []
+def _load_surya_runtime():
+    """Load Surya once per app process and reuse between OCR runs."""
+    global _SURYA_RUNTIME
+    if _SURYA_RUNTIME is not None:
+        return _SURYA_RUNTIME, None, True
 
-    # Try to import/load Surya (new predictor API first, then legacy API)
     surya = None
     surya_error = None
-    log.append(log_line("Loading Surya models (may take a moment)..."))
+
+    # New API (surya-ocr>=0.17 style)
     try:
         from surya.foundation import FoundationPredictor
         from surya.detection import DetectionPredictor
@@ -760,6 +746,9 @@ def run_ocr() -> tuple:
         }
     except Exception as e:
         surya_error = e
+
+    # Legacy API (surya-ocr<=0.6 style)
+    if surya is None:
         try:
             from surya.ocr import run_ocr as surya_run
             from surya.model.detection.model import load_model as load_det
@@ -777,114 +766,239 @@ def run_ocr() -> tuple:
         except Exception as legacy_error:
             surya_error = f"{surya_error}; legacy={legacy_error}"
 
+    if surya is not None:
+        _SURYA_RUNTIME = surya
+        return surya, None, False
+
+    return None, str(surya_error), False
+
+
+def _run_surya_batch(images, surya: dict):
+    """Run a batch of PIL images through Surya using either API shape."""
+    if surya.get("api") == "predictor-v2":
+        return surya["rec_predictor"](
+            images,
+            task_names=[surya["task_name"]] * len(images),
+            det_predictor=surya["det_predictor"],
+            highres_images=images,
+            math_mode=True,
+        )
+
+    return surya["run"](
+        images,
+        [["en"]] * len(images),
+        surya["det_model"],
+        surya["det_proc"],
+        surya["rec_model"],
+        surya["rec_proc"],
+    )
+
+
+def _regions_from_page_result(page_result):
+    regions = []
+    for line in getattr(page_result, "text_lines", []):
+        txt = (getattr(line, "text", "") or "").strip()
+        if not txt:
+            continue
+
+        conf = float(getattr(line, "confidence", 1.0))
+        bbox = getattr(line, "bbox", None)
+        if bbox is None:
+            bbox = getattr(line, "polygon", None)
+
+        regions.append({
+            "text": txt,
+            "confidence": round(conf, 4),
+            "bbox": bbox,
+            "word_count": len(txt.split()),
+        })
+
+    weighted_total = sum(r["confidence"] * r["word_count"] for r in regions)
+    weighted_words = sum(r["word_count"] for r in regions)
+    conf = round(weighted_total / max(weighted_words, 1), 4) if regions else 0.0
+    return regions, conf
+
+
+def run_ocr(progress=gr.Progress(track_tqdm=False)) -> tuple:
+    """Run Surya OCR on all profiled PDFs."""
+    df = load_manifest_df()
+    debug = f"[DEBUG] SS_ROOT={SS_ROOT}\nMANIFEST_CSV={MANIFEST_CSV}\nCSV exists={MANIFEST_CSV.exists()}\n"
+    if not df.empty:
+        debug += f"Manifest rows={len(df)}\nStatuses={df['status'].value_counts().to_dict()}\n"
+    else:
+        debug += "Manifest is EMPTY\n"
+
+    if df.empty:
+        return _ocr_status_html(), debug + "No sources. Complete Steps 1-2 first."
+
+    eligible = df[df["status"].isin(["profiled", "ocred", "rendered"])]
+    if eligible.empty:
+        unknown_or_excluded = df[df["rights_class"].isin(["unknown", "excluded"])] if "rights_class" in df.columns else df.iloc[0:0]
+        pending = df[df["status"] == "pending"] if "status" in df.columns else df.iloc[0:0]
+        debug += (
+            f"Eligible rows={len(eligible)}\n"
+            f"Pending rows={len(pending)}\n"
+            f"Unknown/excluded rights={len(unknown_or_excluded)}\n"
+            "Tip: In Step 1, update rights_class away from unknown/excluded, "
+            "then rerun Step 2 Profile.\n"
+        )
+        return _ocr_status_html(), debug + "No profiled PDFs. Complete Step 2 first."
+
+    log = []
+    queue_rows = []
+
+    progress(0, desc="Preparing OCR run...")
+
+    surya, surya_error, reused = _load_surya_runtime()
     if surya:
-        log.append(log_line(f"✓ Surya models loaded ({surya['api']})"))
+        if reused:
+            log.append(log_line(f"✓ Reusing Surya models ({surya['api']})"))
+        else:
+            log.append(log_line(f"✓ Surya models loaded ({surya['api']})"))
     else:
         log.append(log_line(f"⚠ Surya unavailable ({surya_error}) — falling back to text extraction only"))
 
-    queue_rows = []
+    total_pages_planned = 0
+    for _, row in eligible.iterrows():
+        profile_path = PROFILES_DIR / f"{row['book_id']}_page_profile.json"
+        if not profile_path.exists():
+            continue
+        try:
+            with open(profile_path, encoding="utf-8") as f:
+                profile = json.load(f)
+            total_pages_planned += len(profile.get("pages", []))
+        except Exception:
+            continue
+    total_pages_planned = max(total_pages_planned, 1)
+    processed_pages = 0
 
     for _, row in eligible.iterrows():
-        book_id  = row["book_id"]
+        book_id = row["book_id"]
         profile_path = PROFILES_DIR / f"{book_id}_page_profile.json"
         if not profile_path.exists():
             log.append(log_line(f"✗ {book_id}: no profile"))
             continue
 
-        profile  = json.load(open(profile_path))
+        with open(profile_path, encoding="utf-8") as f:
+            profile = json.load(f)
+
         ocr_pages = []
         review_pages = []
         quarantine_pages = []
         cal = load_calibration()
+        default_cal = cal.get("_default", DEFAULT_CALIBRATION["_default"])
+
+        doc = None
+        try:
+            import fitz
+            pdf_path = SOURCE_DIR / row["filename"]
+            doc = fitz.open(str(pdf_path))
+        except Exception as e:
+            log.append(log_line(f"⚠ {book_id}: PDF open failed ({e})"))
 
-        for page_data in profile["pages"]:
+        # First pass: render OCR/hybrid pages once and keep PIL images for batch OCR.
+        ocr_targets = []
+        for page_data in profile.get("pages", []):
             page_num = page_data["page_number"]
-            route    = page_data["route"]
+            route = page_data["route"]
 
             if route == "embedded_text":
-                # Extract directly via fitz
-                try:
-                    import fitz
-                    pdf_path = SOURCE_DIR / row["filename"]
-                    doc  = fitz.open(str(pdf_path))
-                    page = doc[page_num - 1]
-                    text = page.get_text("text").strip()
-                    doc.close()
-                    regions = [{"text": text, "confidence": 0.99, "bbox": [0,0,100,100], "word_count": len(text.split())}]
-                    conf    = 0.99
-                    method  = "embedded-text"
-                except Exception:
-                    regions = []
-                    conf    = 0.0
-                    method  = "error"
-            elif surya:
+                continue
+
+            render_path = None
+            img = None
+            if doc is not None:
                 try:
-                    from PIL import Image
-                    import fitz
-                    pdf_path = SOURCE_DIR / row["filename"]
-                    doc  = fitz.open(str(pdf_path))
                     page = doc[page_num - 1]
-                    mat  = fitz.Matrix(300/72, 300/72)
-                    pix  = page.get_pixmap(matrix=mat, alpha=False)
-
-                    # Save render
+                    pix = page.get_pixmap(dpi=300, alpha=False)
                     render_dir = RENDERS_DIR / book_id
                     render_dir.mkdir(exist_ok=True)
                     render_path = render_dir / f"{book_id}_page_{page_num:04d}_300dpi.png"
                     pix.save(str(render_path))
-                    doc.close()
+                    page_data["render_path"] = str(render_path.relative_to(SS_ROOT))
+                    page_data["render_dpi"] = 300
 
-                    img = Image.open(render_path).convert("RGB")
-                    if surya["api"] == "predictor-v2":
-                        result = surya["rec_predictor"](
-                            [img],
-                            task_names=[surya["task_name"]],
-                            det_predictor=surya["det_predictor"],
-                            highres_images=[img],
-                            math_mode=True,
-                        )
-                    else:
-                        result = surya["run"](
-                            [img],
-                            [["en"]],
-                            surya["det_model"],
-                            surya["det_proc"],
-                            surya["rec_model"],
-                            surya["rec_proc"],
-                        )
-                    page_result = result[0]
+                    if surya is not None:
+                        from PIL import Image
+                        img = Image.open(render_path).convert("RGB")
+                except Exception as e:
+                    log.append(log_line(f"  ⚠ {book_id} p{page_num}: render failed ({e})"))
 
-                    regions = []
-                    for line in page_result.text_lines:
-                        txt = line.text.strip()
-                        if txt:
-                            c = float(line.confidence) if hasattr(line, "confidence") else 1.0
-                            bbox = getattr(line, "bbox", None)
-                            if bbox is None:
-                                bbox = getattr(line, "polygon", None)
-                            regions.append({"text": txt, "confidence": round(c,4), "bbox": bbox, "word_count": len(txt.split())})
-
-                    conf = sum(r["confidence"]*r["word_count"] for r in regions) / max(sum(r["word_count"] for r in regions), 1) if regions else 0.0
-                    conf = round(conf, 4)
-                    method = "surya"
-
-                    # Update profile with render path
-                    page_data["render_path"] = str(render_path.relative_to(SS_ROOT))
-                    page_data["render_dpi"]  = 300
+            ocr_targets.append({
+                "page_num": page_num,
+                "route": route,
+                "image": img,
+                "render_path": page_data.get("render_path"),
+            })
+
+        # Batch OCR for non-embedded pages.
+        ocr_lookup = {}
+        if surya is not None and ocr_targets:
+            batch_size = SS_SURYA_BATCH_SIZE
+            for start in range(0, len(ocr_targets), batch_size):
+                batch = ocr_targets[start:start + batch_size]
+                batch_pages = [item["page_num"] for item in batch]
+                batch_images = [item["image"] for item in batch if item["image"] is not None]
 
+                try:
+                    if len(batch_images) != len(batch):
+                        raise RuntimeError("One or more page renders missing for OCR batch.")
+
+                    predictions = _run_surya_batch(batch_images, surya)
+                    for item, page_result in zip(batch, predictions):
+                        regions, conf = _regions_from_page_result(page_result)
+                        ocr_lookup[item["page_num"]] = {
+                            "regions": regions,
+                            "confidence": conf,
+                            "method": "surya",
+                        }
                 except Exception as e:
-                    log.append(log_line(f"  ⚠ Page {page_num}: {e}"))
+                    for item in batch:
+                        ocr_lookup[item["page_num"]] = {
+                            "regions": [],
+                            "confidence": 0.0,
+                            "method": "error",
+                        }
+                    log.append(log_line(f"  ⚠ {book_id} batch {batch_pages[0]}-{batch_pages[-1]}: {e}"))
+                finally:
+                    for img in batch_images:
+                        try:
+                            img.close()
+                        except Exception:
+                            pass
+
+        # Second pass: build page-level OCR output and review queue.
+        for page_data in profile.get("pages", []):
+            page_num = page_data["page_number"]
+            route = page_data["route"]
+
+            progress(
+                (processed_pages, total_pages_planned),
+                desc=f"OCR {book_id} p{page_num}/{len(profile.get('pages', []))}",
+            )
+
+            if route == "embedded_text":
+                try:
+                    if doc is None:
+                        raise RuntimeError("PDF document not open")
+                    text = doc[page_num - 1].get_text("text").strip()
+                    regions = [{"text": text, "confidence": 0.99, "bbox": [0, 0, 100, 100], "word_count": len(text.split())}]
+                    conf = 0.99
+                    method = "embedded-text"
+                except Exception:
                     regions = []
-                    conf    = 0.0
-                    method  = "error"
+                    conf = 0.0
+                    method = "error"
+            elif surya:
+                page_out = ocr_lookup.get(page_num, {"regions": [], "confidence": 0.0, "method": "error"})
+                regions = page_out["regions"]
+                conf = page_out["confidence"]
+                method = page_out["method"]
             else:
-                # No surya — skip OCR pages
                 regions = []
-                conf    = 0.0
-                method  = "skipped-no-surya"
+                conf = 0.0
+                method = "skipped-no-surya"
 
-            # Classify confidence
-            default_cal = cal.get("_default", DEFAULT_CALIBRATION["_default"])
             if conf >= default_cal["auto_accept"]:
                 conf_class = "auto-accept"
             elif conf >= default_cal["review"]:
@@ -908,27 +1022,39 @@ def run_ocr() -> tuple:
                 "ocred_at": datetime.utcnow().isoformat() + "Z",
             })
 
-            # Add to review queue
             if conf_class in ("review-required", "low-confidence", "quarantine"):
                 raw_text = " ".join(r["text"] for r in regions)[:500]
                 region_id = f"{book_id}_p{page_num:04d}"
                 queue_rows.append({
-                    "book_id": book_id, "filename": row["filename"],
-                    "page": page_num, "region_id": region_id,
+                    "book_id": book_id,
+                    "filename": row["filename"],
+                    "page": page_num,
+                    "region_id": region_id,
                     "region_class": "narration",
-                    "crop_path": page_data.get("render_path",""),
+                    "crop_path": page_data.get("render_path", ""),
                     "raw_ocr": raw_text,
-                    "confidence": conf, "confidence_class": conf_class,
+                    "confidence": conf,
+                    "confidence_class": conf_class,
                     "status": "quarantine" if conf_class == "quarantine" else "pending",
-                    "reviewer": "", "correction": "", "reason_code": "",
+                    "reviewer": "",
+                    "correction": "",
+                    "reason_code": "",
                 })
 
-        # Save OCR raw
+            processed_pages += 1
+
+        if doc is not None:
+            try:
+                doc.close()
+            except Exception:
+                pass
+
         ocr_dir = OCR_RAW_DIR / book_id
         ocr_dir.mkdir(exist_ok=True)
-        with open(ocr_dir / f"{book_id}_ocr_raw.json", "w") as f:
+        with open(ocr_dir / f"{book_id}_ocr_raw.json", "w", encoding="utf-8") as f:
             json.dump({
-                "book_id": book_id, "filename": row["filename"],
+                "book_id": book_id,
+                "filename": row["filename"],
                 "source_hash": row["sha256"],
                 "page_count": len(ocr_pages),
                 "review_queue": review_pages,
@@ -936,15 +1062,13 @@ def run_ocr() -> tuple:
                 "pages": ocr_pages,
             }, f, indent=2)
 
-        # Save updated profile (with render paths)
-        with open(profile_path, "w") as f:
+        with open(profile_path, "w", encoding="utf-8") as f:
             json.dump(profile, f, indent=2)
 
         df.loc[df["book_id"] == book_id, "status"] = "ocred"
         total_review = len(review_pages) + len(quarantine_pages)
         log.append(log_line(f"✓ {book_id}: {len(ocr_pages)}pp — review queue: {total_review}"))
 
-    # Write review queue
     if queue_rows:
         qdf = pd.DataFrame(queue_rows)
         if QUEUE_CSV.exists():
@@ -952,6 +1076,7 @@ def run_ocr() -> tuple:
             qdf = pd.concat([existing, qdf], ignore_index=True).drop_duplicates(subset=["region_id"])
         qdf.to_csv(QUEUE_CSV, index=False)
 
+    progress((total_pages_planned, total_pages_planned), desc="OCR complete")
     save_manifest_df(df)
     return _ocr_status_html(), "\n".join(log)