Update ocr_helper.py

ocr_helper.py

@@ -2,7 +2,6 @@ import io
 import json
 import logging
 import os
-import glob
 from typing import List, Dict
 from PIL import Image, ImageOps
 import pytesseract
@@ -19,24 +18,33 @@ print(f"===== Application Startup at {__import__('datetime').datetime.now()} ===
 print(f"FVCORE_CACHE: {os.environ.get('FVCORE_CACHE')}")
 print(f"TORCH_HOME: {os.environ.get('TORCH_HOME')}")
 
-# Clean up any files with ?dl=1 in the name BEFORE importing layoutparser
-def cleanup_query_params():
-    """Remove ?dl=1 from filenames in cache directory"""
-    cache_dir = os.environ.get('FVCORE_CACHE', '/root/.torch/iopath_cache')
-    if os.path.exists(cache_dir):
-        for root, dirs, files in os.walk(cache_dir):
-            for filename in files:
-                if '?dl=' in filename:
-                    old_path = os.path.join(root, filename)
-                    new_filename = filename.split('?')[0]
-                    new_path = os.path.join(root, new_filename)
-                    try:
-                        os.rename(old_path, new_path)
-                        logger.info(f"Renamed: {filename} -> {new_filename}")
-                    except Exception as e:
-                        logger.warning(f"Failed to rename {filename}: {e}")
+# Monkey-patch to handle query parameters in filenames
+def patch_file_loading():
+    """Patch torch.load and os.path.isfile to handle ?dl=1 in filenames"""
+    import torch
+    _original_load = torch.load
+    _original_isfile = os.path.isfile
+    
+    def patched_load(f, *args, **kwargs):
+        if isinstance(f, str) and '?dl=' in f:
+            clean_path = f.split('?')[0]
+            if os.path.exists(clean_path):
+                logger.info(f"Using cleaned path: {clean_path}")
+                return _original_load(clean_path, *args, **kwargs)
+        return _original_load(f, *args, **kwargs)
+    
+    def patched_isfile(path):
+        if isinstance(path, str) and '?dl=' in path:
+            clean_path = path.split('?')[0]
+            if _original_isfile(clean_path):
+                return True
+        return _original_isfile(path)
+    
+    torch.load = patched_load
+    os.path.isfile = patched_isfile
+    logger.info("✓ File loading patched to handle query parameters")
 
-cleanup_query_params()
+patch_file_loading()
 
 # ────────────────────────────────────────────────
 # Try importing layoutparser + Detectron2
@@ -48,6 +56,24 @@ try:
     import layoutparser as lp
     logger.info("Layoutparser imported successfully")
     
+    # Also need to patch fvcore's checkpoint loader
+    try:
+        from fvcore.common.checkpoint import Checkpointer
+        _original_load_file = Checkpointer._load_file
+        
+        def patched_load_file(self, f):
+            if isinstance(f, str) and '?dl=' in f:
+                clean_path = f.split('?')[0]
+                if os.path.exists(clean_path):
+                    logger.info(f"Checkpointer using cleaned path: {clean_path}")
+                    return _original_load_file(self, clean_path)
+            return _original_load_file(self, f)
+        
+        Checkpointer._load_file = patched_load_file
+        logger.info("✓ Checkpointer patched")
+    except Exception as e:
+        logger.warning(f"Could not patch Checkpointer: {e}")
+    
     try:
         logger.info("Loading Detectron2 PubLayNet model...")
         # Load Detectron2 model (PubLayNet)
@@ -60,21 +86,8 @@ try:
         logger.info("✅ Detectron2 LayoutParser model loaded successfully.")
     except Exception as e:
         logger.warning(f"⚠️ Failed to load Detectron2 model: {e}")
-        # Try to clean up files and retry once
-        cleanup_query_params()
-        try:
-            logger.info("Retrying model load after cleanup...")
-            TABLE_MODEL = lp.Detectron2LayoutModel(
-                "lp://PubLayNet/faster_rcnn_R_50_FPN_3x/config",
-                extra_config=["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.5],
-                label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"},
-            )
-            USE_LAYOUT = True
-            logger.info("✅ Detectron2 model loaded successfully after cleanup.")
-        except Exception as e2:
-            logger.warning(f"⚠️ Failed to load model after retry: {e2}")
-            TABLE_MODEL = None
-            USE_LAYOUT = False
+        TABLE_MODEL = None
+        USE_LAYOUT = False
 except Exception as e:
     logger.warning(f"⚠️ LayoutParser not available ({e}). Using OCR only.")
     lp = None