Blame for pdf_parser.py

# pdf_parser.py
import os
import re
try:
    import pymupdf as fitz
except ImportError:
    import fitz
from pdf2image import convert_from_path
from PIL import Image
import numpy as np
import cv2
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed

import config

# Configure logging
logger = logging.getLogger(__name__)
logging.basicConfig(level=config.LOG_LEVEL, format=config.LOG_FORMAT)


def convert_pdf_page_to_image(pdf_path, page_num, dpi=config.PDF_TO_IMAGE_DPI):
    """Converts a single page of a PDF to a PIL Image."""
    try:
        images = convert_from_path(pdf_path, dpi=dpi, first_page=page_num + 1, last_page=page_num + 1)
        if images:
            return images[0]
    except Exception as e:
        logger.error(f"Error converting page {page_num} of PDF '{pdf_path}' to image: {e}")
    return None


def pil_to_cv2(pil_image):
    """Converts a PIL Image to an OpenCV image (BGR format)."""
    return cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)


def _extract_native_page_text(doc, page_num, source_pdf):
    """
    Attempts native text extraction from a PDF page using PyMuPDF.
    Includes a safety net to verify text sufficiency and quality.
    
    Returns:
        tuple: (success (bool), page_data (dict or None))
    """
    try:
        page = doc[page_num]
        raw_text = page.get_text("text") or ""
        words = raw_text.split()
        
        # Safety Net Checks:
        # 1. Minimum word count threshold
        if len(words) < config.NATIVE_TEXT_MIN_WORDS:
            return False, None
            
        # 2. Printable character ratio (detect broken/garbled font mappings)
        printable_count = sum(1 for c in raw_text if c.isprintable() or c.isspace())
        if len(raw_text) > 0 and (printable_count / len(raw_text)) < 0.80:
            logger.info(f"Page {page_num + 1} of '{source_pdf}' failed printable character ratio ({printable_count}/{len(raw_text)}). Falling back to OCR.")
            return False, None
            
        # Extract structured blocks from PyMuPDF
        # format: (x0, y0, x1, y1, "lines of text", block_no, block_type)
        blocks = page.get_text("blocks")
        content_items = []
        
        # Check if PyMuPDF table extraction is available
        tables = []
        try:
            if hasattr(page, "find_tables"):
                tab_finder = page.find_tables()
                if tab_finder and tab_finder.tables:
                    for tab in tab_finder.tables:
                        tab_bbox = [int(v) for v in tab.bbox]
                        extracted_rows = tab.extract()
                        flat_cells = [str(c).strip() for row in extracted_rows for c in row if c is not None and str(c).strip()]
                        if flat_cells:
                            tables.append({
                                "bbox": tab_bbox,
                                "cells": flat_cells
                            })
        except Exception as tab_e:
            logger.debug(f"PyMuPDF find_tables not available or failed: {tab_e}")

        for b in blocks:
            if b[6] == 0:  # Text block
                b_text = b[4].strip()
                if not b_text:
                    continue
                x0, y0, x1, y1 = int(b[0]), int(b[1]), int(b[2]), int(b[3])
                
                # Check if this text is already inside an extracted table
                inside_table = False
                for tbl in tables:
                    tx0, ty0, tx1, ty1 = tbl["bbox"]
                    if x0 >= tx0 - 5 and y0 >= ty0 - 5 and x1 <= tx1 + 5 and y1 <= ty1 + 5:
                        inside_table = True
                        break
                        
                if not inside_table:
                    content_items.append((
                        y0,
                        {
                            "type": "plain_text",
                            "text": b_text,
                            "bbox_pil": [x0, y0, x1, y1]
                        }
                    ))
                    
        for tbl in tables:
            x0, y0, x1, y1 = tbl["bbox"]
            is_at_top = (y0 <= 100)
            is_at_bottom = (y1 >= page.rect.height - 100)
            content_items.append((
                y0,
                {
                    "type": "table",
                    "bbox_pil": [x0, y0, x1, y1],
                    "extracted_text_list": tbl["cells"],
                    "raw_cells": [{"text": c, "bbox_pil": [x0, y0, x1, y1]} for c in tbl["cells"]],
                    "is_at_page_top": is_at_top,
                    "is_at_page_bottom": is_at_bottom
                }
            ))

        content_items.sort(key=lambda x: x[0])
        page_data = {
            "page_number": page_num + 1,
            "source_pdf": source_pdf,
            "content": [item[1] for item in content_items]
        }
        return True, page_data

    except Exception as e:
        logger.warning(f"Error during native text extraction on page {page_num + 1} of '{source_pdf}': {e}. Falling back to OCR.")
        return False, None


def _process_page_ocr(pdf_path, page_num, table_detector_pipeline, ocr_reader):
    """Helper function to process a single page using Table Transformer + EasyOCR."""
    page_content_blocks = []
    current_page_data = {"page_number": page_num + 1, "source_pdf": os.path.basename(pdf_path), "content": []}

    try:
        pil_image = convert_pdf_page_to_image(pdf_path, page_num, dpi=config.PDF_TO_IMAGE_DPI)
        if not pil_image:
            logger.warning(f"Could not convert page {page_num} of '{pdf_path}' to image. Skipping page.")
            return current_page_data

        table_boxes_pil = []
        if table_detector_pipeline:
            table_detections = table_detector_pipeline(pil_image)
            for detection in table_detections:
                if detection['label'] == 'table':
                    box = detection['box']
                    x0 = max(0, int(box['xmin']))
                    y0 = max(0, int(box['ymin']))
                    x1 = min(pil_image.width, int(box['xmax']))
                    y1 = min(pil_image.height, int(box['ymax']))
                    if x1 > x0 and y1 > y0:
                        table_boxes_pil.append((x0, y0, x1, y1))

        img_np_rgb = np.array(pil_image)
        non_table_mask = np.ones(img_np_rgb.shape[:2], dtype=np.uint8) * 255
        for x0, y0, x1, y1 in table_boxes_pil:
            non_table_mask[y0:y1, x0:x1] = 0
        
        non_table_img_np = cv2.bitwise_and(img_np_rgb, img_np_rgb, mask=non_table_mask)
        non_table_ocr_results = ocr_reader.readtext(non_table_img_np, paragraph=True) if ocr_reader else []

        for ocr_result in non_table_ocr_results:
            if len(ocr_result) == 2:
                bbox, text = ocr_result
            elif len(ocr_result) == 3:
                bbox, text, _ = ocr_result
            else:
                continue

            if bbox and isinstance(bbox, list) and len(bbox) > 0 and \
               isinstance(bbox[0], (list, tuple)) and len(bbox[0]) == 2:
                pos_y = int(bbox[0][1])
                page_content_blocks.append(
                    (pos_y, {"type": "plain_text", "text": text, "bbox_pil": [int(c) for pt in bbox for c in pt]})
                )

        page_height = pil_image.height
        tolerance = 100  # Pixels tolerance for checking if table is at edge

        for x0_tbl, y0_tbl, x1_tbl, y1_tbl in table_boxes_pil:
            table_pil_image_crop = pil_image.crop((x0_tbl, y0_tbl, x1_tbl, y1_tbl))
            table_ocr_results = ocr_reader.readtext(np.array(table_pil_image_crop)) if ocr_reader else []

            table_cells_text = []
            raw_table_cells = []
            for item in table_ocr_results:
                if len(item) == 3:
                    bbox_cell, text_cell, prob_cell = item
                elif len(item) == 2:
                    bbox_cell, text_cell = item
                else:
                    continue
                table_cells_text.append(text_cell)
                adjusted_bbox_cell = [[int(pt[0] + x0_tbl), int(pt[1] + y0_tbl)] for pt in bbox_cell]
                raw_table_cells.append({"text": text_cell, "bbox_pil": [int(c) for pt in adjusted_bbox_cell for c in pt]})

            if table_cells_text:
                is_at_top = (y0_tbl <= tolerance)
                is_at_bottom = (y1_tbl >= page_height - tolerance)
                page_content_blocks.append(
                    (y0_tbl, {
                        "type": "table",
                        "bbox_pil": [x0_tbl, y0_tbl, x1_tbl, y1_tbl],
                        "extracted_text_list": table_cells_text,
                        "raw_cells": raw_table_cells,
                        "is_at_page_top": is_at_top,
                        "is_at_page_bottom": is_at_bottom
                    })
                )
        
        page_content_blocks.sort(key=lambda x: x[0])
        current_page_data["content"] = [block[1] for block in page_content_blocks]

    except Exception as e:
        logger.error(f"Error OCR processing page {page_num} of PDF '{pdf_path}': {e}", exc_info=True)
    
    return current_page_data


def _process_single_page_hybrid(pdf_path, page_num, doc, table_detector_pipeline, ocr_reader):
    """
    Executes native PyMuPDF text extraction first with safety net;
    falls back to OCR if word density is low or page is scanned.
    """
    source_pdf = os.path.basename(pdf_path)
    
    if getattr(config, "USE_NATIVE_PDF_TEXT", True) and doc is not None:
        success, native_data = _extract_native_page_text(doc, page_num, source_pdf)
        if success and native_data and native_data.get("content"):
            logger.debug(f"Native extraction successful for page {page_num + 1} of '{source_pdf}'.")
            return native_data
            
    logger.info(f"Using OCR fallback for page {page_num + 1} of '{source_pdf}'.")
    return _process_page_ocr(pdf_path, page_num, table_detector_pipeline, ocr_reader)


def extract_content_from_pdf(pdf_path, table_detector_pipeline=None, ocr_reader=None, max_workers=4):
    """
    Extracts structured content (text and tables) from a PDF file.
    Uses native text extraction where available and falls back to OCR.

    Args:
        pdf_path (str): Path to the PDF file.
        table_detector_pipeline: Initialized table detection pipeline (optional for native-only).
        ocr_reader: Initialized EasyOCR reader (optional for native-only).
        max_workers (int): Maximum number of threads for parallel processing.

    Returns:
        list: A list of dicts with 'page_number' and 'content' list.
    """
    extracted_pdf_data = []
    try:
        doc = fitz.open(pdf_path)
        num_pages = len(doc)
    except Exception as e:
        logger.error(f"Error opening PDF '{pdf_path}': {e}")
        return extracted_pdf_data

    # Check how many pages can be natively parsed
    all_pages_native = True
    page_results = [None] * num_pages
    
    if getattr(config, "USE_NATIVE_PDF_TEXT", True):
        for page_num in range(num_pages):
            success, native_data = _extract_native_page_text(doc, page_num, os.path.basename(pdf_path))
            if success and native_data and native_data.get("content"):
                page_results[page_num] = native_data
            else:
                all_pages_native = False

    if all_pages_native and all(p is not None for p in page_results):
        doc.close()
        logger.info(f"All {num_pages} pages of '{os.path.basename(pdf_path)}' extracted natively via PyMuPDF.")
        return sorted(page_results, key=lambda x: x['page_number'])

    # For pages requiring OCR, run OCR in thread pool
    missing_page_nums = [i for i, r in enumerate(page_results) if r is None]
    logger.info(f"PDF '{os.path.basename(pdf_path)}': {len(page_results) - len(missing_page_nums)}/{num_pages} pages parsed natively. Running OCR fallback on {len(missing_page_nums)} pages.")

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {
            executor.submit(_process_page_ocr, pdf_path, page_num, table_detector_pipeline, ocr_reader): page_num 
            for page_num in missing_page_nums
        }
        for future in as_completed(futures):
            p_num = futures[future]
            page_results[p_num] = future.result()

    doc.close()
    extracted_pdf_data = [p for p in page_results if p is not None]
    extracted_pdf_data = sorted(extracted_pdf_data, key=lambda x: x['page_number'])
    return extracted_pdf_data