Update working_yolo_pipeline.py

working_yolo_pipeline.py

@@ -1514,6 +1514,241 @@ def post_process_json_with_inference(json_data, classifier):
 
 #     return final_output, page_separator_x
 
+#==========================================================================================================================================================================================
+
+# def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
+#                             page_num: int, fitz_page: fitz.Page,
+#                             pdf_name: str) -> Tuple[List[Dict[str, Any]], Optional[int]]:
+#     """
+#     OPTIMIZED FLOW - MODIFIED FOR CORRECT ORDERING:
+#     1. Run YOLO to find Equations/Tables.
+#     2. Store detections with page_num but DON'T assign global IDs yet
+#     3. Mask raw text with YOLO boxes.
+#     4. Run Column Detection on the MASKED data.
+#     5. Proceed with OCR (Native or High-Res Tesseract Fallback) and Output.
+#     """
+#     # NOTE: Removed global counter increments from here
+
+#     start_time_total = time.time()
+
+#     if original_img is None:
+#         print(f"  ❌ Invalid image for page {page_num}.")
+#         return None, None
+
+#     # ====================================================================
+#     # --- 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=CONF_THRESHOLD, verbose=False)
+
+#     relevant_detections = []
+#     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) ---
+#     # ====================================================================
+#     raw_words_for_layout = get_word_data_for_detection(
+#         fitz_page, pdf_path, page_num,
+#         top_margin_percent=0.10, bottom_margin_percent=0.10
+#     )
+
+#     masked_word_data = merge_yolo_into_word_data(raw_words_for_layout, merged_detections, scale_factor=2.0)
+
+#     # ====================================================================
+#     # --- STEP 3: COLUMN DETECTION ---
+#     # ====================================================================
+#     page_width_pdf = fitz_page.rect.width
+#     page_height_pdf = fitz_page.rect.height
+
+#     column_detection_params = {
+#         'cluster_bin_size': 2, 'cluster_smoothing': 2,
+#         'cluster_min_width': 10, 'cluster_threshold_percentile': 85,
+#     }
+
+#     separators = calculate_x_gutters(masked_word_data, column_detection_params, page_height_pdf)
+
+#     page_separator_x = None
+#     if separators:
+#         central_min = page_width_pdf * 0.35
+#         central_max = page_width_pdf * 0.65
+#         central_separators = [s for s in separators if central_min <= s <= central_max]
+
+#         if central_separators:
+#             center_x = page_width_pdf / 2
+#             page_separator_x = min(central_separators, key=lambda x: abs(x - center_x))
+#             print(f"      ✅ Column Split Confirmed at X={page_separator_x:.1f}")
+#         else:
+#             print("      ⚠️ Gutter found off-center. Ignoring.")
+#     else:
+#         print("      -> Single Column Layout Confirmed.")
+
+#     # ====================================================================
+#     # --- STEP 4: COMPONENT EXTRACTION (MODIFIED - Store without ID) ---
+#     # ====================================================================
+#     start_time_components = time.time()
+#     component_metadata = []
+
+#     for detection in merged_detections:
+#         x1, y1, x2, y2 = detection['coords']
+#         class_name = detection['class']
+
+#         # DON'T assign global IDs here - just store the type and coordinates
+#         component_crop = original_img[y1:y2, x1:x2]
+        
+#         # Store image temporarily with page and position info in filename
+#         temp_filename = f"{pdf_name}_page{page_num}_{class_name}_y{y1}.png"
+#         temp_filepath = os.path.join(FIGURE_EXTRACTION_DIR, temp_filename)
+#         cv2.imwrite(temp_filepath, component_crop)
+
+#         y_midpoint = (y1 + y2) // 2
+#         component_metadata.append({
+#             'type': class_name, 
+#             'word': f"TEMP_{class_name.upper()}_PAGE{page_num}_Y{y1}",  # Temporary placeholder
+#             'bbox': [int(x1), int(y1), int(x2), int(y2)],
+#             'y0': int(y_midpoint), 
+#             'x0': int(x1),
+#             'page_num': page_num,  # CRITICAL: Store page number
+#             'temp_filepath': temp_filepath  # Store temp filepath for later renaming
+#         })
+
+#     # ====================================================================
+#     # --- STEP 5: HYBRID OCR (Native Text + Cached Tesseract Fallback) ---
+#     # ====================================================================
+#     raw_ocr_output = []
+#     scale_factor = 2.0
+
+#     try:
+#         raw_ocr_output = extract_native_words_and_convert(fitz_page, scale_factor=scale_factor)
+#     except Exception as e:
+#         print(f"  ❌ Native text extraction failed: {e}")
+
+#     if not raw_ocr_output:
+#         if _ocr_cache.has_ocr(pdf_path, page_num):
+#             print(f"  ⚡ Using cached Tesseract OCR for page {page_num}")
+#             cached_word_data = _ocr_cache.get_ocr(pdf_path, page_num)
+#             for word_tuple in cached_word_data:
+#                 word_text, x1, y1, x2, y2 = word_tuple
+#                 x1_pix = int(x1 * scale_factor)
+#                 y1_pix = int(y1 * scale_factor)
+#                 x2_pix = int(x2 * scale_factor)
+#                 y2_pix = int(y2 * scale_factor)
+
+#                 raw_ocr_output.append({
+#                     'type': 'text', 'word': word_text, 'confidence': 95.0,
+#                     'bbox': [x1_pix, y1_pix, x2_pix, y2_pix],
+#                     'y0': y1_pix, 'x0': x1_pix
+#                 })
+#         else:
+#             try:
+#                 ocr_zoom = 4.0
+#                 pix_ocr = fitz_page.get_pixmap(matrix=fitz.Matrix(ocr_zoom, ocr_zoom))
+#                 img_ocr_np = np.frombuffer(pix_ocr.samples, dtype=np.uint8).reshape(pix_ocr.height, pix_ocr.width,
+#                                                                                     pix_ocr.n)
+#                 if pix_ocr.n == 3:
+#                     img_ocr_np = cv2.cvtColor(img_ocr_np, cv2.COLOR_RGB2BGR)
+#                 elif pix_ocr.n == 4:
+#                     img_ocr_np = cv2.cvtColor(img_ocr_np, cv2.COLOR_RGBA2BGR)
+
+#                 processed_img = preprocess_image_for_ocr(img_ocr_np)
+#                 custom_config = r'--oem 3 --psm 6'
+#                 hocr_data = pytesseract.image_to_data(
+#                     processed_img,
+#                     output_type=pytesseract.Output.DICT,
+#                     config=custom_config
+#                 )
+
+#                 for i in range(len(hocr_data['level'])):
+#                     text = hocr_data['text'][i]
+#                     cleaned_text = sanitize_text(text).strip()
+                    
+#                     if cleaned_text and hocr_data['conf'][i] > -1:
+#                         scale_adjustment = scale_factor / ocr_zoom
+#                         x1 = int(hocr_data['left'][i] * scale_adjustment)
+#                         y1 = int(hocr_data['top'][i] * scale_adjustment)
+#                         w = int(hocr_data['width'][i] * scale_adjustment)
+#                         h = int(hocr_data['height'][i] * scale_adjustment)
+#                         x2 = x1 + w
+#                         y2 = y1 + h
+
+#                         raw_ocr_output.append({
+#                             'type': 'text',
+#                             'word': cleaned_text,
+#                             'confidence': float(hocr_data['conf'][i]),
+#                             'bbox': [x1, y1, x2, y2],
+#                             'y0': y1,
+#                             'x0': x1
+#                         })
+#             except Exception as e:
+#                 print(f"  ❌ Tesseract OCR Error: {e}")
+    
+#     # ====================================================================
+#     # --- STEP 6: OCR CLEANING AND MERGING ---
+#     # ====================================================================
+#     items_to_sort = []
+
+#     for ocr_word in raw_ocr_output:
+#         is_suppressed = False
+#         for component in component_metadata:
+#             ioa = calculate_ioa(ocr_word['bbox'], component['bbox'])
+#             if ioa > IOA_SUPPRESSION_THRESHOLD:
+#                 is_suppressed = True
+#                 break
+#         if not is_suppressed:
+#             items_to_sort.append(ocr_word)
+
+#     items_to_sort.extend(component_metadata)
+
+#     # ====================================================================
+#     # --- STEP 7: LINE-BASED SORTING ---
+#     # ====================================================================
+#     items_to_sort.sort(key=lambda x: (x['y0'], x['x0']))
+#     lines = []
+
+#     for item in items_to_sort:
+#         placed = False
+#         for line in lines:
+#             y_ref = min(it['y0'] for it in line)
+#             if abs(y_ref - item['y0']) < LINE_TOLERANCE:
+#                 line.append(item)
+#                 placed = True
+#                 break
+#         if not placed and item['type'] in ['equation', 'figure']:
+#             for line in lines:
+#                 y_ref = min(it['y0'] for it in line)
+#                 if abs(y_ref - item['y0']) < 20:
+#                     line.append(item)
+#                     placed = True
+#                     break
+#         if not placed:
+#             lines.append([item])
+
+#     for line in lines:
+#         line.sort(key=lambda x: x['x0'])
+
+#     final_output = []
+#     for line in lines:
+#         for item in line:
+#             data_item = {"word": item["word"], "bbox": item["bbox"], "type": item["type"]}
+#             if 'tag' in item: data_item['tag'] = item['tag']
+#             if 'page_num' in item: data_item['page_num'] = item['page_num']
+#             if 'temp_filepath' in item: data_item['temp_filepath'] = item['temp_filepath']
+#             final_output.append(data_item)
+
+#     return final_output, page_separator_x
+#=================================================================================================================================================================================================
+
 
 
 def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
@@ -1543,17 +1778,22 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
     results = model.predict(source=original_img, conf=CONF_THRESHOLD, verbose=False)
 
     relevant_detections = []
+    
+    # FIX 1: Use .data.tolist() to preserve float coordinates (matches feedback.py)
     if results and results[0].boxes:
-        for box in results[0].boxes:
-            class_id = int(box.cls[0])
-            class_name = model.names[class_id]
+        for box in results[0].boxes.data.tolist():
+            x1, y1, x2, y2, conf, cls_id = box
+            class_name = model.names[int(cls_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])}
+                    {'coords': (x1, y1, x2, y2), 'y1': y1, 'class': class_name, 'conf': conf}
                 )
 
     merged_detections = merge_overlapping_boxes(relevant_detections, IOU_MERGE_THRESHOLD)
+    
+    # FIX 2: Add the missing filter_nested_boxes step (matches feedback.py)
+    merged_detections = filter_nested_boxes(merged_detections, IOA_SUPPRESSION_THRESHOLD)
+    
     print(f"    [LOG] YOLO found {len(merged_detections)} objects in {time.time() - start_time_yolo:.3f}s.")
 
     # ====================================================================
@@ -1601,12 +1841,18 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
     component_metadata = []
 
     for detection in merged_detections:
-        x1, y1, x2, y2 = detection['coords']
+        # FIX 3: Cast float coordinates to int HERE for numpy array slicing
+        x1, y1, x2, y2 = map(int, detection['coords'])
         class_name = detection['class']
 
+        # Ensure coordinates are within image bounds
+        h, w = original_img.shape[:2]
+        x1, y1 = max(0, x1), max(0, y1)
+        x2, y2 = min(w, x2), min(h, y2)
+        
         # DON'T assign global IDs here - just store the type and coordinates
         component_crop = original_img[y1:y2, x1:x2]
-        
+
         # Store image temporarily with page and position info in filename
         temp_filename = f"{pdf_name}_page{page_num}_{class_name}_y{y1}.png"
         temp_filepath = os.path.join(FIGURE_EXTRACTION_DIR, temp_filename)
@@ -1614,10 +1860,10 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
 
         y_midpoint = (y1 + y2) // 2
         component_metadata.append({
-            'type': class_name, 
+            'type': class_name,
             'word': f"TEMP_{class_name.upper()}_PAGE{page_num}_Y{y1}",  # Temporary placeholder
-            'bbox': [int(x1), int(y1), int(x2), int(y2)],
-            'y0': int(y_midpoint), 
+            'bbox': [x1, y1, x2, y2],
+            'y0': int(y_midpoint),
             'x0': int(x1),
             'page_num': page_num,  # CRITICAL: Store page number
             'temp_filepath': temp_filepath  # Store temp filepath for later renaming
@@ -1672,7 +1918,7 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
                 for i in range(len(hocr_data['level'])):
                     text = hocr_data['text'][i]
                     cleaned_text = sanitize_text(text).strip()
-                    
+
                     if cleaned_text and hocr_data['conf'][i] > -1:
                         scale_adjustment = scale_factor / ocr_zoom
                         x1 = int(hocr_data['left'][i] * scale_adjustment)
@@ -1692,7 +1938,7 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
                         })
             except Exception as e:
                 print(f"  ❌ Tesseract OCR Error: {e}")
-    
+
     # ====================================================================
     # --- STEP 6: OCR CLEANING AND MERGING ---
     # ====================================================================
@@ -1750,11 +1996,6 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
 
 
 
-
-
-
-
-
 def run_single_pdf_preprocessing(pdf_path: str, preprocessed_json_path: str) -> Optional[str]:
     global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT