Optimize pipeline: load models once per PDF instead of per page for 10x speed improvement

pipeline.py

@@ -31,6 +31,10 @@ from qr.qr_extraction import process_image_no_save as process_qr
 from signature.inference import detect_signatures
 from stamp_detector.detect import detect_stamps_no_save
 
+# Import for model loading
+from ultralytics import YOLO
+import os
+
 
 def pdf_to_images(pdf_path: str, dpi: int = 200) -> List[np.ndarray]:
     """
@@ -68,13 +72,56 @@ def pdf_to_images(pdf_path: str, dpi: int = 200) -> List[np.ndarray]:
     return images
 
 
+def _load_signature_model(signature_model_path: Optional[str] = None):
+    """Load signature model once for reuse."""
+    from huggingface_hub import hf_hub_download
+    
+    if signature_model_path and Path(signature_model_path).exists():
+        model_path = signature_model_path
+    else:
+        local_model_path = Path("yolov8s.pt")
+        if local_model_path.exists():
+            model_path = str(local_model_path)
+        else:
+            try:
+                hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
+                model_path = hf_hub_download(
+                    repo_id="tech4humans/yolov8s-signature-detector",
+                    filename="yolov8s.pt",
+                    token=hf_token
+                )
+            except Exception as e:
+                raise RuntimeError(f"Failed to load signature model: {e}")
+    
+    print("📥 Loading signature model...")
+    model = YOLO(model_path)
+    print("✓ Signature model loaded")
+    return model
+
+
+def _load_stamp_model(stamp_model_path: str = "stamp_detector/stamp_model.pt"):
+    """Load stamp model once for reuse."""
+    if not Path(stamp_model_path).exists():
+        default_path = Path("stamp_detector/stamp_model.pt")
+        if default_path.exists():
+            stamp_model_path = str(default_path)
+        else:
+            raise FileNotFoundError(f"Stamp model not found: {stamp_model_path}")
+    
+    print("📥 Loading stamp model...")
+    model = YOLO(stamp_model_path)
+    print("✓ Stamp model loaded")
+    return model
+
+
 def process_pdf_pipeline(
     pdf_path: str,
     output_dir: str = "pipeline_outputs",
     stamp_model_path: str = "stamp_detector/stamp_model.pt",
     stamp_conf: float = 0.25,
     dpi: int = 200,
-    save_intermediate: bool = False
+    save_intermediate: bool = False,
+    signature_model_path: Optional[str] = None
 ) -> Dict[str, Any]:
     """
     Process a PDF file by converting each page to an image and running the pipeline.
@@ -86,6 +133,7 @@ def process_pdf_pipeline(
         stamp_conf: Confidence threshold for stamp detection
         dpi: DPI for PDF to image conversion
         save_intermediate: Whether to save intermediate results
+        signature_model_path: Path to signature model (optional, will auto-download if not provided)
         
     Returns:
         Combined results dictionary for all pages
@@ -104,8 +152,17 @@ def process_pdf_pipeline(
     print(f"Processing PDF: {pdf_path.name}")
     print(f"{'='*70}\n")
     
+    # Load models once before processing pages
+    print("🔄 Loading models (this happens once for all pages)...")
+    try:
+        signature_model = _load_signature_model(signature_model_path)
+        stamp_model = _load_stamp_model(stamp_model_path)
+    except Exception as e:
+        print(f"✗ Error loading models: {str(e)}")
+        raise
+    
     # Convert PDF to images
-    print(f"📄 Converting PDF pages to images (DPI: {dpi})...")
+    print(f"\n📄 Converting PDF pages to images (DPI: {dpi})...")
     try:
         page_images = pdf_to_images(str(pdf_path), dpi=dpi)
         print(f"✓ Converted {len(page_images)} page(s) to images\n")
@@ -126,12 +183,13 @@ def process_pdf_pipeline(
             temp_img_path = temp_dir / f"page_{page_num}.jpg"
             cv2.imwrite(str(temp_img_path), img)
             
-            # Process the page
+            # Process the page with pre-loaded models
             try:
                 page_result = process_image_pipeline(
                     str(temp_img_path),
                     output_dir=output_dir,
-                    stamp_model_path=stamp_model_path,
+                    signature_model=signature_model,
+                    stamp_model=stamp_model,
                     stamp_conf=stamp_conf,
                     save_intermediate=save_intermediate
                 )
@@ -189,7 +247,9 @@ def process_image_pipeline(
     signature_model_path: Optional[str] = None,
     stamp_model_path: str = "stamp_detector/stamp_model.pt",
     stamp_conf: float = 0.25,
-    save_intermediate: bool = False
+    save_intermediate: bool = False,
+    signature_model: Optional[Any] = None,
+    stamp_model: Optional[Any] = None
 ) -> Dict[str, Any]:
     """
     Process a single image through all three detection models.
@@ -254,9 +314,16 @@ def process_image_pipeline(
     print(f"\n🔷 Step 2/3: Signature Detection")
     print("-" * 70)
     try:
+        # Use pre-loaded model if provided, otherwise load on demand
+        if signature_model is None:
+            if signature_model_path:
+                signature_model = _load_signature_model(signature_model_path)
+            else:
+                signature_model = _load_signature_model()
+        
         sig_result = detect_signatures(
             str(image_path),
-            model=None,  # Will auto-load
+            model=signature_model,  # Use pre-loaded model
             output_dir=None,  # Don't save
             signatures_dir=None,  # Don't save
             save_crops=False  # Don't save crops
@@ -284,13 +351,17 @@ def process_image_pipeline(
     print(f"\n🔷 Step 3/3: Stamp Detection")
     print("-" * 70)
     try:
-        if not Path(stamp_model_path).exists():
-            raise FileNotFoundError(f"Stamp model not found: {stamp_model_path}")
+        # Use pre-loaded model if provided, otherwise load on demand
+        if stamp_model is None:
+            if not Path(stamp_model_path).exists():
+                raise FileNotFoundError(f"Stamp model not found: {stamp_model_path}")
+            stamp_model = _load_stamp_model(stamp_model_path)
         
         stamp_result = detect_stamps_no_save(
             str(image_path),
             model_path=stamp_model_path,
-            conf=stamp_conf
+            conf=stamp_conf,
+            model=stamp_model  # Pass pre-loaded model
         )
         
         if stamp_result and stamp_result.get("detections"):

stamp_detector/detect.py

@@ -9,7 +9,7 @@ import json
 from ultralytics import YOLO
 
 
-def detect_stamps_no_save(image_path, model_path="stamp_model.pt", conf=0.25):
+def detect_stamps_no_save(image_path, model_path="stamp_model.pt", conf=0.25, model=None):
     """
     Detect stamps without saving images.
     
@@ -17,27 +17,30 @@ def detect_stamps_no_save(image_path, model_path="stamp_model.pt", conf=0.25):
         image_path: Path to input image
         model_path: Path to model (or will download from HF Hub if not found)
         conf: Confidence threshold
+        model: Pre-loaded YOLO model (optional, will load if not provided)
         
     Returns:
         dict: Detection results with detections and image_size
     """
-    # Load model - try to download from HF Hub if not found locally
-    if not os.path.exists(model_path):
-        # Try to download from Hugging Face Hub
-        try:
-            from huggingface_hub import hf_hub_download
-            print(f"Model not found locally, attempting to download from HF Hub...")
-            # You can upload your model to HF Hub and use it here
-            # For now, try the default path in stamp_detector directory
-            default_path = os.path.join("stamp_detector", "stamp_model.pt")
-            if os.path.exists(default_path):
-                model_path = default_path
-            else:
-                raise FileNotFoundError(f"Stamp model not found: {model_path}. Please upload stamp_model.pt to the Space.")
-        except ImportError:
-            raise FileNotFoundError(f"Stamp model not found: {model_path}")
-    
-    model = YOLO(model_path)
+    # Use pre-loaded model if provided, otherwise load model
+    if model is None:
+        # Load model - try to download from HF Hub if not found locally
+        if not os.path.exists(model_path):
+            # Try to download from Hugging Face Hub
+            try:
+                from huggingface_hub import hf_hub_download
+                print(f"Model not found locally, attempting to download from HF Hub...")
+                # You can upload your model to HF Hub and use it here
+                # For now, try the default path in stamp_detector directory
+                default_path = os.path.join("stamp_detector", "stamp_model.pt")
+                if os.path.exists(default_path):
+                    model_path = default_path
+                else:
+                    raise FileNotFoundError(f"Stamp model not found: {model_path}. Please upload stamp_model.pt to the Space.")
+            except ImportError:
+                raise FileNotFoundError(f"Stamp model not found: {model_path}")
+        
+        model = YOLO(model_path)
     
     # Load image
     if not os.path.exists(image_path):