Update working_yolo_pipeline.py

working_yolo_pipeline.py

@@ -184,7 +184,7 @@ FIGURE_EXTRACTION_DIR = './figure_extraction'
 TEMP_IMAGE_DIR = './temp_pdf_images'
 
 # Detection parameters
-CONF_THRESHOLD = 0.2
+# CONF_THRESHOLD = 0.2
 TARGET_CLASSES = ['figure', 'equation']
 IOU_MERGE_THRESHOLD = 0.4
 IOA_SUPPRESSION_THRESHOLD = 0.7
@@ -1094,22 +1094,55 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
     # --- STEP 1: YOLO DETECTION ---
     # ====================================================================
     start_time_yolo = time.time()
-    results = model.predict(source=original_img, conf=CONF_THRESHOLD, imgsz=640, verbose=False)
+    results = model.predict(source=original_img, conf=0.2, imgsz=640, verbose=False)
+
+    
+
 
     relevant_detections = []
+
+    THRESHOLDS = {
+        'figure': 0.75,
+        'equation': 0.20
+    }
+
+
+
+
     if results and results[0].boxes:
         for box in results[0].boxes:
             class_id = int(box.cls[0])
             class_name = model.names[class_id]
-            if class_name in TARGET_CLASSES:
-                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
-                relevant_detections.append(
-                    {'coords': (x1, y1, x2, y2), 'y1': y1, 'class': class_name, 'conf': float(box.conf[0])}
-                )
+            conf = float(box.conf[0])
+
+            # Logic: Check if class is in our list AND meets its specific threshold
+            if class_name in THRESHOLDS:
+                if conf >= THRESHOLDS[class_name]:
+                    x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
+                    relevant_detections.append({
+                        'coords': (x1, y1, x2, y2), 
+                        'y1': y1, 
+                        'class': class_name, 
+                        'conf': conf
+                    })
 
     merged_detections = merge_overlapping_boxes(relevant_detections, IOU_MERGE_THRESHOLD)
     print(f"    [LOG] YOLO found {len(merged_detections)} objects in {time.time() - start_time_yolo:.3f}s.")
 
+    
+    # if results and results[0].boxes:
+    #     for box in results[0].boxes:
+    #         class_id = int(box.cls[0])
+    #         class_name = model.names[class_id]
+    #         if class_name in TARGET_CLASSES:
+    #             x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
+    #             relevant_detections.append(
+    #                 {'coords': (x1, y1, x2, y2), 'y1': y1, 'class': class_name, 'conf': float(box.conf[0])}
+    #             )
+
+    # merged_detections = merge_overlapping_boxes(relevant_detections, IOU_MERGE_THRESHOLD)
+    # print(f"    [LOG] YOLO found {len(merged_detections)} objects in {time.time() - start_time_yolo:.3f}s.")
+
     # ====================================================================
     # --- STEP 2: PREPARE DATA FOR COLUMN DETECTION (MASKING) ---
     # ====================================================================