test/verify_bbox.py

#!/usr/bin/env python3
"""Visual bbox verification tool for Stage A output.

This script renders a PDF with colored bounding boxes overlaid to verify
that the Stage A parser correctly identifies and locates elements.

Color coding:
- Blue: FLOWING_TEXT (regular text blocks)
- Green: IN_PLACE (headers, footers, captions)
- Red: BYPASS (figures, pictures)
- Purple: TABLE (table boundaries)
- Orange: EQUATION (formulas)
- Yellow (thin): TABLE cells

Usage:
    python scripts/verify_bbox.py --input sample.pdf --output verify_output.pdf
    python scripts/verify_bbox.py --input sample.pdf --output verify_output.pdf --pages 0,1,2
"""

import argparse
import logging
import sys
from pathlib import Path

import fitz  # PyMuPDF

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))

from pdf2zh.parser import StageAParser
from pdf2zh.parser.enums import ElementCategory
from pdf2zh.parser.schema import validate_stage_output

logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
logger = logging.getLogger(__name__)

# Color definitions (RGB tuples, 0-1 scale)
CATEGORY_COLORS = {
    ElementCategory.FLOWING_TEXT: (0.0, 0.0, 1.0),  # Blue
    ElementCategory.IN_PLACE: (0.0, 0.8, 0.0),  # Green
    ElementCategory.BYPASS: (1.0, 0.0, 0.0),  # Red
    ElementCategory.TABLE: (0.5, 0.0, 0.5),  # Purple
    ElementCategory.EQUATION: (1.0, 0.5, 0.0),  # Orange
}

CELL_COLOR = (0.8, 0.8, 0.0)  # Yellow for table cells
TEXT_CELL_COLOR = (1.0, 0.0, 1.0)


def draw_bbox(
    page: fitz.Page,
    bbox: list[float],
    img_width: float,
    img_height: float,
    color: tuple,
    width: float = 2.0,
) -> None:

    page_rect = page.rect
    pdf_width = page_rect.width
    pdf_height = page_rect.height

    scale_x = pdf_width / img_width
    scale_y = pdf_height / img_height

    x0 = page_rect.x0 + (bbox[0] * scale_x)
    y0 = page_rect.y0 + (bbox[1] * scale_y)
    x1 = page_rect.x0 + (bbox[2] * scale_x)
    y1 = page_rect.y0 + (bbox[3] * scale_y)

    rect = fitz.Rect(x0, y0, x1, y1)

    page.draw_rect(rect, color=color, width=width)


def draw_label(page: fitz.Page, bbox: list[float], label: str, color: tuple) -> None:
    """Draw a label above the bbox.

    Args:
        page: fitz Page to draw on
        bbox: [x0, y0, x1, y1] in PDF points
        label: Text label to display
        color: RGB tuple for text color
    """
    # Position label above the bbox
    text_point = fitz.Point(bbox[0], bbox[1] - 2)

    # Draw label with small font
    page.insert_text(
        text_point,
        label,
        fontsize=8,
        color=color,
    )


def verify_pdf(
    input_path: str,
    output_path: str,
    pages: list[int] | None = None,
    device: str = "auto",
) -> None:
    """Parse a PDF and create a verification output with bbox overlays.

    Args:
        input_path: Path to input PDF
        output_path: Path to save verification PDF
        pages: Optional list of page indices to process
        device: Device for Surya models
    """
    input_path = Path(input_path)
    output_path = Path(output_path)

    if not input_path.exists():
        raise FileNotFoundError(f"Input PDF not found: {input_path}")

    logger.info(f"Parsing {input_path}...")

    # Parse the PDF through explicit Stage A phases
    parser = StageAParser(device=device)
    parsed_doc = parser.parse_pdf(input_path, pages=pages)

    logger.info(f"Found {len(parsed_doc.pages)} pages")

    # Open the original PDF
    doc = fitz.open(input_path)

    # Draw bboxes on each page
    for page_data in parsed_doc.pages:
        page_idx = page_data.page_index
        if page_idx >= len(doc):
            continue

        page = doc[page_idx]
        logger.info(f"Page {page_idx}: {len(page_data.elements)} elements")

        # Draw element bboxes
        for elem in page_data.elements:
            color = CATEGORY_COLORS.get(elem.category, (0.5, 0.5, 0.5))
            draw_bbox(
                page,
                elem.bbox_pdf,
                page_data.page_width,
                page_data.page_height,
                color,
                width=2.0,
            )
            draw_label(page, elem.bbox_pdf, f"{elem.label}", color)

            # Draw cell bboxes for tables
            if elem.category == ElementCategory.TABLE:
                for cell in elem.cells:
                    draw_bbox(
                        page,
                        cell.bbox_pdf,
                        page_data.page_width,
                        page_data.page_height,
                        CELL_COLOR,
                        width=1.5,
                    )
                    draw_bbox(
                        page,
                        cell.bbox_text,
                        page_data.page_width,
                        page_data.page_height,
                        TEXT_CELL_COLOR,
                        width=1.0,
                    )

    # Save the annotated PDF
    output_path.parent.mkdir(parents=True, exist_ok=True)
    doc.save(output_path)
    doc.close()

    logger.info(f"Saved verification PDF to {output_path}")

    # Save JSON output alongside the PDF
    json_path = output_path.with_suffix(".json")
    json_path.write_text(parsed_doc.to_json(indent=2), encoding="utf-8")
    logger.info(f"Saved JSON to {json_path}")

    # Run schema validation
    validation = validate_stage_output(
        parsed_doc.to_dict(), stage="A", skip_json_schema=True
    )

    # Print summary
    print("\nVerification Summary:")
    print("=" * 50)
    print(f"Input:  {input_path}")
    print(f"PDF:    {output_path}")
    print(f"JSON:   {json_path}")
    print(f"Pages:  {len(parsed_doc.pages)}")

    total_elements = sum(len(p.elements) for p in parsed_doc.pages)
    print(f"Elements: {total_elements}")

    # Count by category
    category_counts = {}
    for page_data in parsed_doc.pages:
        for elem in page_data.elements:
            cat = elem.category.value
            category_counts[cat] = category_counts.get(cat, 0) + 1

    print("\nElements by category:")
    for cat, count in sorted(category_counts.items()):
        print(f"  {cat}: {count}")

    # Validation result
    if validation.valid:
        print("\nSchema validation: PASS")
    else:
        print(f"\nSchema validation: FAIL ({len(validation.errors)} errors)")
        for err in validation.errors:
            print(f"  [{err.code}] {err.path}: {err.message}")

    print("\nColor legend:")
    print("  Blue: FLOWING_TEXT")
    print("  Green: IN_PLACE")
    print("  Red: BYPASS")
    print("  Purple: TABLE")
    print("  Orange: EQUATION")
    print("  Yellow (thin): Table cells")


def main():
    parser = argparse.ArgumentParser(
        description="Verify Stage A bbox detection with visual output",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python scripts/verify_bbox.py --input sample.pdf --output verify.pdf
    python scripts/verify_bbox.py --input sample.pdf --output verify.pdf --pages 0,1,2
    python scripts/verify_bbox.py --input sample.pdf --output verify.pdf --device cpu
        """,
    )

    parser.add_argument("--input", "-i", required=True, help="Input PDF file path")
    parser.add_argument(
        "--output", "-o", required=True, help="Output verification PDF path"
    )
    parser.add_argument(
        "--pages",
        "-p",
        type=str,
        default=None,
        help="Comma-separated list of page indices (0-based)",
    )
    parser.add_argument(
        "--device",
        "-d",
        type=str,
        default="auto",
        choices=["auto", "cuda", "mps", "cpu"],
        help="Device for Surya models (default: auto)",
    )

    args = parser.parse_args()

    # Parse pages if specified
    pages = None
    if args.pages:
        pages = [int(p.strip()) for p in args.pages.split(",")]

    try:
        verify_pdf(
            input_path=args.input,
            output_path=args.output,
            pages=pages,
            device=args.device,
        )
    except FileNotFoundError as e:
        logger.error(str(e))
        sys.exit(1)
    except Exception as e:
        logger.exception(f"Error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()