Prevent OCR stalls by timing out Surya load and falling back

smoke_signal_tab.py

@@ -26,6 +26,7 @@ import hashlib
 import json
 import os
 import tempfile
+import threading
 import time
 from datetime import datetime
 from pathlib import Path
@@ -709,10 +710,17 @@ def _profile_status_html() -> str:
 
 # ── Step 3: OCR ────────────────────────────────────────────────────────────────
 _SURYA_RUNTIME = None
+_SURYA_LOAD_THREAD = None
+_SURYA_LOAD_ERROR = None
+_SURYA_LOAD_LOCK = threading.Lock()
 try:
     SS_SURYA_BATCH_SIZE = max(1, int(os.environ.get("SS_SURYA_BATCH_SIZE", "4")))
 except Exception:
     SS_SURYA_BATCH_SIZE = 4
+try:
+    SS_SURYA_LOAD_TIMEOUT_SEC = max(1, int(os.environ.get("SS_SURYA_LOAD_TIMEOUT_SEC", "20")))
+except Exception:
+    SS_SURYA_LOAD_TIMEOUT_SEC = 20
 
 
 def _load_surya_runtime():
@@ -773,6 +781,47 @@ def _load_surya_runtime():
     return None, str(surya_error), False
 
 
+def _surya_loader_worker():
+    """Background loader to avoid blocking OCR forever on slow model downloads."""
+    global _SURYA_LOAD_ERROR
+    surya, err, _ = _load_surya_runtime()
+    if surya is None:
+        _SURYA_LOAD_ERROR = err
+    else:
+        _SURYA_LOAD_ERROR = None
+
+
+def _get_surya_runtime_with_timeout(timeout_sec: int):
+    """
+    Return Surya runtime quickly.
+    If model load is still in progress after timeout, caller should fallback this run.
+    """
+    global _SURYA_LOAD_THREAD
+
+    if _SURYA_RUNTIME is not None:
+        return _SURYA_RUNTIME, None, True
+
+    with _SURYA_LOAD_LOCK:
+        if _SURYA_RUNTIME is not None:
+            return _SURYA_RUNTIME, None, True
+
+        if _SURYA_LOAD_THREAD is None or not _SURYA_LOAD_THREAD.is_alive():
+            _SURYA_LOAD_THREAD = threading.Thread(target=_surya_loader_worker, daemon=True)
+            _SURYA_LOAD_THREAD.start()
+
+        loader_thread = _SURYA_LOAD_THREAD
+
+    loader_thread.join(timeout=timeout_sec)
+
+    if _SURYA_RUNTIME is not None:
+        return _SURYA_RUNTIME, None, False
+
+    if loader_thread.is_alive():
+        return None, f"Surya load exceeded {timeout_sec}s (still loading in background)", False
+
+    return None, _SURYA_LOAD_ERROR or "Surya load failed", 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":
@@ -819,6 +868,67 @@ def _regions_from_page_result(page_result):
     return regions, conf
 
 
+def _run_tesseract_batch(images):
+    """
+    Tesseract fallback for OCR when Surya is unavailable/slow.
+    Returns list of dicts with regions/confidence/method aligned to input order.
+    """
+    try:
+        import pytesseract
+    except Exception as e:
+        return [{"regions": [], "confidence": 0.0, "method": f"error-no-tesseract ({e})"} for _ in images]
+
+    outputs = []
+    for img in images:
+        try:
+            data = pytesseract.image_to_data(
+                img,
+                lang="eng",
+                config="--oem 1 --psm 6",
+                output_type=pytesseract.Output.DICT,
+            )
+            regions = []
+            conf_weighted = 0.0
+            word_count = 0
+            n = len(data.get("text", []))
+
+            for i in range(n):
+                txt = str(data["text"][i]).strip()
+                if not txt:
+                    continue
+
+                try:
+                    conf_raw = float(data["conf"][i])
+                except Exception:
+                    conf_raw = -1.0
+                if conf_raw < 0:
+                    continue
+
+                conf = max(0.0, min(1.0, conf_raw / 100.0))
+                left = int(data["left"][i])
+                top = int(data["top"][i])
+                width = int(data["width"][i])
+                height = int(data["height"][i])
+                bbox = [left, top, left + width, top + height]
+                wc = max(len(txt.split()), 1)
+
+                regions.append({
+                    "text": txt,
+                    "confidence": round(conf, 4),
+                    "bbox": bbox,
+                    "word_count": wc,
+                })
+                conf_weighted += conf * wc
+                word_count += wc
+
+            avg_conf = round(conf_weighted / max(word_count, 1), 4) if regions else 0.0
+            outputs.append({"regions": regions, "confidence": avg_conf, "method": "tesseract"})
+        except Exception as e:
+            outputs.append({"regions": [], "confidence": 0.0, "method": f"error-tesseract ({e})"})
+
+    return outputs
+
+
 def run_ocr(progress=gr.Progress(track_tqdm=False)) -> tuple:
     """Run Surya OCR on all profiled PDFs."""
     df = load_manifest_df()
@@ -849,14 +959,15 @@ def run_ocr(progress=gr.Progress(track_tqdm=False)) -> tuple:
 
     progress(0, desc="Preparing OCR run...")
 
-    surya, surya_error, reused = _load_surya_runtime()
+    progress(0, desc="Loading OCR engine...")
+    surya, surya_error, reused = _get_surya_runtime_with_timeout(SS_SURYA_LOAD_TIMEOUT_SEC)
     if surya:
         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"))
+        log.append(log_line(f"⚠ Surya unavailable ({surya_error}) — using Tesseract fallback for this run"))
 
     total_pages_planned = 0
     for _, row in eligible.iterrows():
@@ -933,7 +1044,7 @@ def run_ocr(progress=gr.Progress(track_tqdm=False)) -> tuple:
 
         # Batch OCR for non-embedded pages.
         ocr_lookup = {}
-        if surya is not None and ocr_targets:
+        if 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]
@@ -944,22 +1055,34 @@ def run_ocr(progress=gr.Progress(track_tqdm=False)) -> tuple:
                     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",
-                        }
+                    if surya is not None:
+                        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",
+                            }
+                    else:
+                        fallback_preds = _run_tesseract_batch(batch_images)
+                        for item, pred in zip(batch, fallback_preds):
+                            ocr_lookup[item["page_num"]] = pred
                 except Exception as 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}"))
+                    # If Surya batch fails, try Tesseract for this batch before giving up.
+                    if surya is not None:
+                        log.append(log_line(f"  ⚠ {book_id} batch {batch_pages[0]}-{batch_pages[-1]} Surya error: {e}; retrying with Tesseract"))
+                        fallback_preds = _run_tesseract_batch(batch_images)
+                        for item, pred in zip(batch, fallback_preds):
+                            ocr_lookup[item["page_num"]] = pred
+                    else:
+                        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:
@@ -989,7 +1112,7 @@ def run_ocr(progress=gr.Progress(track_tqdm=False)) -> tuple:
                     regions = []
                     conf = 0.0
                     method = "error"
-            elif surya:
+            elif surya or page_num in ocr_lookup:
                 page_out = ocr_lookup.get(page_num, {"regions": [], "confidence": 0.0, "method": "error"})
                 regions = page_out["regions"]
                 conf = page_out["confidence"]