Fix Tesseract fallback batch loading and OCR progress visibility

smoke_signal_tab.py

@@ -721,6 +721,14 @@ 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
+try:
+    SS_RENDER_DPI_SURYA = max(72, int(os.environ.get("SS_RENDER_DPI_SURYA", "300")))
+except Exception:
+    SS_RENDER_DPI_SURYA = 300
+try:
+    SS_RENDER_DPI_FALLBACK = max(72, int(os.environ.get("SS_RENDER_DPI_FALLBACK", "220")))
+except Exception:
+    SS_RENDER_DPI_FALLBACK = 220
 
 
 def _load_surya_runtime():
@@ -982,6 +990,8 @@ def run_ocr(progress=gr.Progress(track_tqdm=False)) -> tuple:
             continue
     total_pages_planned = max(total_pages_planned, 1)
     processed_pages = 0
+    rendered_pages = 0
+    total_work_units = max(total_pages_planned * 2, 1)
 
     for _, row in eligible.iterrows():
         book_id = row["book_id"]
@@ -1009,36 +1019,36 @@ def run_ocr(progress=gr.Progress(track_tqdm=False)) -> tuple:
 
         # First pass: render OCR/hybrid pages once and keep PIL images for batch OCR.
         ocr_targets = []
+        render_dpi = SS_RENDER_DPI_SURYA if surya is not None else SS_RENDER_DPI_FALLBACK
         for page_data in profile.get("pages", []):
             page_num = page_data["page_number"]
             route = page_data["route"]
+            rendered_pages += 1
+            progress(
+                (rendered_pages, total_work_units),
+                desc=f"Rendering {book_id} p{page_num} @ {render_dpi}dpi",
+            )
 
             if route == "embedded_text":
                 continue
 
             render_path = None
-            img = None
             if doc is not None:
                 try:
                     page = doc[page_num - 1]
-                    pix = page.get_pixmap(dpi=300, alpha=False)
+                    pix = page.get_pixmap(dpi=render_dpi, 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"
+                    render_path = render_dir / f"{book_id}_page_{page_num:04d}_{render_dpi}dpi.png"
                     pix.save(str(render_path))
                     page_data["render_path"] = str(render_path.relative_to(SS_ROOT))
-                    page_data["render_dpi"] = 300
-
-                    if surya is not None:
-                        from PIL import Image
-                        img = Image.open(render_path).convert("RGB")
+                    page_data["render_dpi"] = render_dpi
                 except Exception as e:
                     log.append(log_line(f"  ⚠ {book_id} p{page_num}: render failed ({e})"))
 
             ocr_targets.append({
                 "page_num": page_num,
                 "route": route,
-                "image": img,
                 "render_path": page_data.get("render_path"),
             })
 
@@ -1049,15 +1059,52 @@ def run_ocr(progress=gr.Progress(track_tqdm=False)) -> tuple:
             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]
+                batch_images = []
+                batch_items_with_images = []
+                try:
+                    from PIL import Image
+                except Exception as e:
+                    Image = None
+                    log.append(log_line(f"  ⚠ PIL unavailable for OCR batch ({e})"))
+
+                if Image is not None:
+                    for item in batch:
+                        rel_path = item.get("render_path")
+                        if not rel_path:
+                            ocr_lookup[item["page_num"]] = {
+                                "regions": [],
+                                "confidence": 0.0,
+                                "method": "error-no-render-path",
+                            }
+                            continue
+
+                        render_abs = SS_ROOT / rel_path
+                        if not render_abs.exists():
+                            ocr_lookup[item["page_num"]] = {
+                                "regions": [],
+                                "confidence": 0.0,
+                                "method": "error-render-missing",
+                            }
+                            continue
+
+                        try:
+                            img = Image.open(render_abs).convert("RGB")
+                            batch_images.append(img)
+                            batch_items_with_images.append(item)
+                        except Exception as e:
+                            ocr_lookup[item["page_num"]] = {
+                                "regions": [],
+                                "confidence": 0.0,
+                                "method": f"error-open-image ({e})",
+                            }
 
                 try:
-                    if len(batch_images) != len(batch):
-                        raise RuntimeError("One or more page renders missing for OCR batch.")
+                    if not batch_items_with_images:
+                        raise RuntimeError("No render images available in this OCR batch.")
 
                     if surya is not None:
                         predictions = _run_surya_batch(batch_images, surya)
-                        for item, page_result in zip(batch, predictions):
+                        for item, page_result in zip(batch_items_with_images, predictions):
                             regions, conf = _regions_from_page_result(page_result)
                             ocr_lookup[item["page_num"]] = {
                                 "regions": regions,
@@ -1066,17 +1113,17 @@ def run_ocr(progress=gr.Progress(track_tqdm=False)) -> tuple:
                             }
                     else:
                         fallback_preds = _run_tesseract_batch(batch_images)
-                        for item, pred in zip(batch, fallback_preds):
+                        for item, pred in zip(batch_items_with_images, fallback_preds):
                             ocr_lookup[item["page_num"]] = pred
                 except Exception as 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):
+                        for item, pred in zip(batch_items_with_images, fallback_preds):
                             ocr_lookup[item["page_num"]] = pred
                     else:
-                        for item in batch:
+                        for item in batch_items_with_images:
                             ocr_lookup[item["page_num"]] = {
                                 "regions": [],
                                 "confidence": 0.0,
@@ -1096,7 +1143,7 @@ def run_ocr(progress=gr.Progress(track_tqdm=False)) -> tuple:
             route = page_data["route"]
 
             progress(
-                (processed_pages, total_pages_planned),
+                (total_pages_planned + processed_pages, total_work_units),
                 desc=f"OCR {book_id} p{page_num}/{len(profile.get('pages', []))}",
             )
 
@@ -1199,7 +1246,7 @@ def run_ocr(progress=gr.Progress(track_tqdm=False)) -> 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")
+    progress((total_work_units, total_work_units), desc="OCR complete")
     save_manifest_df(df)
     return _ocr_status_html(), "\n".join(log)