lambdagent/agentexample/qaagent67lambda/scripts/extract_pdfs_v2.py

#!/usr/bin/env python3
"""
PDF extraction v2: table-aware
- Extracts tables as Markdown tables
- Marks table boundaries with [TABLE_START] [TABLE_END] tags
- Regular text extracted normally
"""
import os
import sys
import re
import fitz
from pathlib import Path

from config import cfg
RAW_DIR = cfg.raw_dir
OUT_DIR = cfg.processed_dir
OUT_DIR.mkdir(exist_ok=True)

pdf_files = sorted(RAW_DIR.glob('*.pdf'))
total = len(pdf_files)
print(f'Total: {total} PDFs')

success = 0
failed = 0
tables_found = 0


def table_to_markdown(table_data, headers=None):
    """Convert extracted table data to Markdown table format"""
    if not table_data or not table_data[0]:
        return ''

    rows = []
    for row in table_data:
        # Clean cells: replace None with empty, remove newlines
        cells = []
        for cell in row:
            if cell is None:
                cells.append('')
            else:
                cells.append(str(cell).replace('\n', ' ').strip())
        rows.append(cells)

    if not rows:
        return ''

    # Ensure all rows have same number of columns
    max_cols = max(len(r) for r in rows)
    for r in rows:
        while len(r) < max_cols:
            r.append('')

    # Build markdown table
    lines = []
    # Header row
    lines.append('| ' + ' | '.join(rows[0]) + ' |')
    lines.append('| ' + ' | '.join(['---'] * max_cols) + ' |')
    # Data rows
    for row in rows[1:]:
        lines.append('| ' + ' | '.join(row) + ' |')

    return '\n'.join(lines)


for i, pdf_path in enumerate(pdf_files):
    out_path = OUT_DIR / (pdf_path.stem + '.txt')

    try:
        doc = fitz.open(str(pdf_path))
        text_parts = []

        for page_num in range(len(doc)):
            page = doc[page_num]

            # Try to find tables on this page
            page_tables = page.find_tables()
            table_rects = []

            if page_tables.tables:
                for tab in page_tables.tables:
                    table_data = tab.extract()
                    if table_data and len(table_data) > 1:
                        md_table = table_to_markdown(table_data)
                        if md_table:
                            text_parts.append(f'[TABLE_START]\n{md_table}\n[TABLE_END]')
                            tables_found += 1
                            table_rects.append(tab.bbox)

            # Get regular text, excluding table areas
            if not table_rects:
                page_text = page.get_text()
                if page_text.strip():
                    text_parts.append(page_text)
            else:
                # Extract text only from non-table regions
                # Get page full rect, then extract text blocks and skip those inside table rects
                blocks = page.get_text("blocks")  # list of (x0, y0, x1, y1, text, block_no, type)
                for block in blocks:
                    bx0, by0, bx1, by1, block_text = block[:5]
                    if not block_text.strip():
                        continue
                    # Check if this text block overlaps any table rect
                    in_table = False
                    for trect in table_rects:
                        tx0, ty0, tx1, ty1 = trect
                        # Check overlap: block center inside table rect
                        bcx = (bx0 + bx1) / 2
                        bcy = (by0 + by1) / 2
                        if tx0 <= bcx <= tx1 and ty0 <= bcy <= ty1:
                            in_table = True
                            break
                    if not in_table:
                        text_parts.append(block_text)

        doc.close()

        full_text = '\n'.join(text_parts)
        if full_text.strip():
            with open(out_path, 'w', encoding='utf-8') as f:
                f.write(f'# {pdf_path.name}\n\n')
                f.write(full_text)
            success += 1
        else:
            with open(out_path, 'w', encoding='utf-8') as f:
                f.write(f'# {pdf_path.name}\n# [scan PDF]\n')
            failed += 1

    except Exception as e:
        failed += 1
        if (i + 1) % 100 == 0:
            print(f'  ERROR [{i+1}/{total}] {pdf_path.name}: {e}', file=sys.stderr)

    if (i + 1) % 200 == 0:
        print(f'  Progress: {i+1}/{total}, tables: {tables_found}')

print(f'\nDone! success={success} failed={failed} tables={tables_found}')