feat(F-09): rich parsing — formula + figure extraction
- codex/parsing/mathpix.py: pix2tex (local, CPU) primary + MathPix API optional; bbox heuristic h>15px, math-char-count>5; singleton model cache - codex/parsing/figures.py: pymupdf embedded-image extraction → PNG; caption detection via proximity + "Figure/Fig./Abbildung" prefix - codex/models.py: FormulaChunk + FigureChunk dataclasses (R-10/R-11) - codex/ingest.py: --rich flag wires formula+figure extraction post-ingest - codex/cli.py: search_app sub-typer (paper + formula subcommands), --rich flag on ingest; wiki_app from F-12 preserved intact - codex/config.py: mathpix_app_id/key, pix2tex_fallback, figures_dir - infra/schema.sql: formulas + figures tables with HNSW pgvector indexes - pyproject.toml: pymupdf>=1.24, pix2tex>=0.1.4 - tests/parsing/test_mathpix.py + test_figures.py: 31 tests (mock pix2tex + MathPix HTTP, real pymupdf on synthetic PDF) Gate: 158 passed, ruff clean, mypy clean (20 files) Requirements: R-10 R-11 R-12 R-13 R-14 → done Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
166
codex/parsing/figures.py
Normal file
166
codex/parsing/figures.py
Normal file
@@ -0,0 +1,166 @@
|
||||
"""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
|
||||
|
||||
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,
|
||||
)
|
||||
)
|
||||
|
||||
return results
|
||||
Reference in New Issue
Block a user