Update working_yolo_pipeline.py

working_yolo_pipeline.py

@@ -73,6 +73,23 @@ except Exception as e:
 
 
 
+from typing import Optional
+
+def sanitize_text(text: Optional[str]) -> str:
+    """Removes surrogate characters and other invalid code points that cause UTF-8 encoding errors."""
+    if not isinstance(text, str) or text is None:
+        return ""
+    
+    # Matches all surrogates (\ud800-\udfff) and common non-characters (\ufffe, \uffff).
+    # This specifically removes '\udefd' which is causing your error.
+    surrogates_and_nonchars = re.compile(r'[\ud800-\udfff\ufffe\uffff]')
+    
+    # Replace the invalid characters with a standard space.
+    # We strip afterward in the calling function.
+    return surrogates_and_nonchars.sub(' ', text)
+
+
+
 
 def get_latex_from_base64(base64_string: str) -> str:
     """
@@ -559,20 +576,45 @@ def pixmap_to_numpy(pix: fitz.Pixmap) -> np.ndarray:
     return img
 
 
+# def extract_native_words_and_convert(fitz_page, scale_factor: float = 2.0) -> list:
+#     raw_word_data = fitz_page.get_text("words")
+#     converted_ocr_output = []
+#     DEFAULT_CONFIDENCE = 99.0
+
+#     for x1, y1, x2, y2, word, *rest in raw_word_data:
+#         if not word.strip(): continue
+#         x1_pix = int(x1 * scale_factor)
+#         y1_pix = int(y1 * scale_factor)
+#         x2_pix = int(x2 * scale_factor)
+#         y2_pix = int(y2 * scale_factor)
+#         converted_ocr_output.append({
+#             'type': 'text',
+#             'word': word,
+#             'confidence': DEFAULT_CONFIDENCE,
+#             'bbox': [x1_pix, y1_pix, x2_pix, y2_pix],
+#             'y0': y1_pix, 'x0': x1_pix
+#         })
+#     return converted_ocr_output
+
+
+
 def extract_native_words_and_convert(fitz_page, scale_factor: float = 2.0) -> list:
     raw_word_data = fitz_page.get_text("words")
     converted_ocr_output = []
     DEFAULT_CONFIDENCE = 99.0
 
     for x1, y1, x2, y2, word, *rest in raw_word_data:
-        if not word.strip(): continue
+        # --- FIX: SANITIZE TEXT HERE ---
+        cleaned_word = sanitize_text(word)
+        if not cleaned_word.strip(): continue
+        
         x1_pix = int(x1 * scale_factor)
         y1_pix = int(y1 * scale_factor)
         x2_pix = int(x2 * scale_factor)
         y2_pix = int(y2 * scale_factor)
         converted_ocr_output.append({
             'type': 'text',
-            'word': word,
+            'word': cleaned_word, # Use the sanitized word
             'confidence': DEFAULT_CONFIDENCE,
             'bbox': [x1_pix, y1_pix, x2_pix, y2_pix],
             'y0': y1_pix, 'x0': x1_pix
@@ -580,6 +622,272 @@ def extract_native_words_and_convert(fitz_page, scale_factor: float = 2.0) -> li
     return converted_ocr_output
 
 
+
+
+    
+
+# 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:
+#     1. Run YOLO to find Equations/Tables.
+#     2. Mask raw text with YOLO boxes.
+#     3. Run Column Detection on the MASKED data.
+#     4. Proceed with OCR (Native or High-Res Tesseract Fallback) and Output.
+#     """
+#     global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT
+
+#     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)
+
+#     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) ---
+#     # ====================================================================
+#     # Note: This uses the updated 'get_word_data_for_detection' which has its own optimizations
+#     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 (Save Images) ---
+#     # ====================================================================
+#     start_time_components = time.time()
+#     component_metadata = []
+#     fig_count_page = 0
+#     eq_count_page = 0
+
+#     for detection in merged_detections:
+#         x1, y1, x2, y2 = detection['coords']
+#         class_name = detection['class']
+
+#         if class_name == 'figure':
+#             GLOBAL_FIGURE_COUNT += 1
+#             counter = GLOBAL_FIGURE_COUNT
+#             component_word = f"FIGURE{counter}"
+#             fig_count_page += 1
+#         elif class_name == 'equation':
+#             GLOBAL_EQUATION_COUNT += 1
+#             counter = GLOBAL_EQUATION_COUNT
+#             component_word = f"EQUATION{counter}"
+#             eq_count_page += 1
+#         else:
+#             continue
+
+#         component_crop = original_img[y1:y2, x1:x2]
+#         component_filename = f"{pdf_name}_page{page_num}_{class_name}{counter}.png"
+#         cv2.imwrite(os.path.join(FIGURE_EXTRACTION_DIR, component_filename), component_crop)
+
+#         y_midpoint = (y1 + y2) // 2
+#         component_metadata.append({
+#             'type': class_name, 'word': component_word,
+#             'bbox': [int(x1), int(y1), int(x2), int(y2)],
+#             'y0': int(y_midpoint), 'x0': int(x1)
+#         })
+
+#     # ====================================================================
+#     # --- STEP 5: HYBRID OCR (Native Text + Cached Tesseract Fallback) ---
+#     # ====================================================================
+#     raw_ocr_output = []
+#     scale_factor = 2.0  # Pipeline standard scale
+
+#     try:
+#         # Try getting native text first
+#         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 native text is missing, fall back to OCR
+#     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
+
+#                 # Scale from PDF points to Pipeline Pixels (2.0)
+#                 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:
+#             # === START OF OPTIMIZED OCR BLOCK ===
+#             try:
+#                 # 1. Re-render Page at High Resolution (Zoom 4.0 = ~300 DPI)
+#                 # We do this specifically for OCR accuracy, separate from the pipeline image
+#                 ocr_zoom = 4.0
+#                 pix_ocr = fitz_page.get_pixmap(matrix=fitz.Matrix(ocr_zoom, ocr_zoom))
+
+#                 # Convert PyMuPDF Pixmap to OpenCV format
+#                 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)
+
+#                 # 2. Preprocess (Binarization)
+#                 # Ensure 'preprocess_image_for_ocr' is defined at top of file!
+#                 processed_img = preprocess_image_for_ocr(img_ocr_np)
+
+#                 # 3. Run Tesseract with Optimized Configuration
+#                 # --oem 3: Default LSTM engine
+#                 # --psm 6: Assume a single uniform block of text (Critical for lists/questions)
+#                 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].strip()
+#                     if text and hocr_data['conf'][i] > -1:
+#                         # 4. Coordinate Mapping
+#                         # We scanned at Zoom 4.0, but our pipeline expects Zoom 2.0.
+#                         # Scale Factor = (Target 2.0) / (Source 4.0) = 0.5
+#                         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': 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}")
+#             # === END OF OPTIMIZED OCR BLOCK ===
+
+#     # ====================================================================
+#     # --- STEP 6: OCR CLEANING AND MERGING ---
+#     # ====================================================================
+#     items_to_sort = []
+
+#     for ocr_word in raw_ocr_output:
+#         is_suppressed = False
+#         for component in component_metadata:
+#             # Do not include words that are inside figure/equation boxes
+#             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)
+
+#     # Add figures/equations back into the flow as "words"
+#     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']
+#             final_output.append(data_item)
+
+#     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]]:
@@ -701,6 +1009,7 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
 
     try:
         # Try getting native text first
+        # NOTE: extract_native_words_and_convert MUST ALSO BE UPDATED TO USE sanitize_text
         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}")
@@ -728,7 +1037,6 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
             # === START OF OPTIMIZED OCR BLOCK ===
             try:
                 # 1. Re-render Page at High Resolution (Zoom 4.0 = ~300 DPI)
-                # We do this specifically for OCR accuracy, separate from the pipeline image
                 ocr_zoom = 4.0
                 pix_ocr = fitz_page.get_pixmap(matrix=fitz.Matrix(ocr_zoom, ocr_zoom))
 
@@ -741,12 +1049,9 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
                     img_ocr_np = cv2.cvtColor(img_ocr_np, cv2.COLOR_RGBA2BGR)
 
                 # 2. Preprocess (Binarization)
-                # Ensure 'preprocess_image_for_ocr' is defined at top of file!
                 processed_img = preprocess_image_for_ocr(img_ocr_np)
 
                 # 3. Run Tesseract with Optimized Configuration
-                # --oem 3: Default LSTM engine
-                # --psm 6: Assume a single uniform block of text (Critical for lists/questions)
                 custom_config = r'--oem 3 --psm 6'
 
                 hocr_data = pytesseract.image_to_data(
@@ -756,11 +1061,13 @@ 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].strip()
-                    if text and hocr_data['conf'][i] > -1:
+                    text = hocr_data['text'][i] # Retrieve raw Tesseract text
+                    
+                    # --- FIX: SANITIZE TEXT AND THEN STRIP ---
+                    cleaned_text = sanitize_text(text).strip()
+                    
+                    if cleaned_text and hocr_data['conf'][i] > -1:
                         # 4. Coordinate Mapping
-                        # We scanned at Zoom 4.0, but our pipeline expects Zoom 2.0.
-                        # Scale Factor = (Target 2.0) / (Source 4.0) = 0.5
                         scale_adjustment = scale_factor / ocr_zoom
 
                         x1 = int(hocr_data['left'][i] * scale_adjustment)
@@ -772,7 +1079,7 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
 
                         raw_ocr_output.append({
                             'type': 'text',
-                            'word': text,
+                            'word': cleaned_text, # Use the sanitized word
                             'confidence': float(hocr_data['conf'][i]),
                             'bbox': [x1, y1, x2, y2],
                             'y0': y1,
@@ -781,7 +1088,7 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
             except Exception as e:
                 print(f"  ❌ Tesseract OCR Error: {e}")
             # === END OF OPTIMIZED OCR BLOCK ===
-
+    
     # ====================================================================
     # --- STEP 6: OCR CLEANING AND MERGING ---
     # ====================================================================