codex-py/codex/parsing/figures.py

"""Figure extraction from PDF pages.

Uses PyMuPDF (``fitz``) to iterate over embedded images in each page,
render them to PNG and detect nearby captions using text-block heuristics.

The public entry point is :func:`extract_figures`.
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any

from codex.models import FigureChunk

logger = logging.getLogger(__name__)

# Caption prefixes (case-insensitive check against text block content).
_CAPTION_PREFIXES: tuple[str, ...] = ("figure", "fig.", "abbildung")

# Maximum vertical distance (in points) from the image bottom or top
# to consider a text block as the caption.
_CAPTION_TOLERANCE_Y = 60.0


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _looks_like_caption(text: str) -> bool:
    """Return True if *text* starts with a known caption prefix."""
    lowered = text.strip().lower()
    return any(lowered.startswith(prefix) for prefix in _CAPTION_PREFIXES)


def _find_caption(
    img_rect: Any,  # fitz.Rect
    text_blocks: list[Any],
) -> str:
    """Find the closest caption text block to *img_rect*.

    Searches text blocks that:
    1. Start with a caption prefix (e.g. "Figure", "Fig.", "Abbildung").
    2. Are within ``_CAPTION_TOLERANCE_Y`` points above or below the image.

    Returns the first matching text, stripped; empty string if none found.
    """
    best: str = ""
    best_dist = float("inf")

    for block in text_blocks:
        if len(block) < 5:
            continue
        bx0, by0, bx1, by1 = block[0], block[1], block[2], block[3]
        text: str = block[4]

        if not _looks_like_caption(text):
            continue

        # Distance: gap below image or gap above image
        gap_below = by0 - img_rect.y1
        gap_above = img_rect.y0 - by1

        if 0 <= gap_below <= _CAPTION_TOLERANCE_Y:
            dist = gap_below
        elif 0 <= gap_above <= _CAPTION_TOLERANCE_Y:
            dist = gap_above
        else:
            continue

        # Horizontal overlap sanity check (caption should overlap with image)
        overlap = min(bx1, img_rect.x1) - max(bx0, img_rect.x0)
        if overlap < 0:
            continue

        if dist < best_dist:
            best_dist = dist
            best = text.replace("\n", " ").strip()

    return best


# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------


def extract_figures(
    pdf_path: str,
    output_dir: str,
) -> list[FigureChunk]:
    """Extract embedded figures from *pdf_path* and save them as PNG files.

    For each image on each page:
    1. Extract the image via ``page.get_images(full=True)``.
    2. Render the image as a PNG via ``fitz.Pixmap``.
    3. Search nearby text blocks for a caption (``Figure …`` / ``Fig. …`` /
       ``Abbildung …``).
    4. Save the PNG to ``<output_dir>/<paper_id>_p<page>_fig<n>.png``.
    5. Yield a :class:`~codex.models.FigureChunk`.

    Parameters
    ----------
    pdf_path:
        Path to the source PDF.
    output_dir:
        Directory where PNG images are written (created if absent).

    Returns
    -------
    list[FigureChunk]
        Extracted figure metadata; empty list if no images found.
    """
    import fitz

    results: list[FigureChunk] = []
    paper_id = Path(pdf_path).stem
    out_dir = Path(output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    try:
        doc: Any = fitz.open(pdf_path)
    except Exception:
        logger.exception("fitz.open failed for %s", pdf_path)
        return results

    try:
        for page_num in range(len(doc)):
            page = doc[page_num]
            images: list[Any] = page.get_images(full=True)
            text_blocks: list[Any] = page.get_text("blocks")

            for fig_idx, img_info in enumerate(images):
                xref: int = img_info[0]

                # Get image bounding box on the page
                img_rects: list[Any] = page.get_image_rects(xref)
                if not img_rects:
                    continue
                img_rect = img_rects[0]

                # Extract and save the image pixmap
                try:
                    pix: Any = fitz.Pixmap(doc, xref)
                    if pix.n > 4:  # CMYK or similar — convert to RGB
                        pix = fitz.Pixmap(fitz.csRGB, pix)
                    filename = f"{paper_id}_p{page_num + 1}_fig{fig_idx + 1}.png"
                    img_path = out_dir / filename
                    pix.save(str(img_path))
                except Exception:
                    logger.debug("Could not extract image xref=%d on page %d", xref, page_num + 1)
                    continue

                caption = _find_caption(img_rect, text_blocks)

                results.append(
                    FigureChunk(
                        paper_id=paper_id,
                        page=page_num + 1,
                        image_path=str(img_path),
                        caption=caption,
                    )
                )
    finally:
        doc.close()

    return results