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:
Tarik Moussa
2026-06-14 01:16:24 +02:00
parent d2f9141a5c
commit 1df9be6563
13 changed files with 1930 additions and 26 deletions

166
codex/parsing/figures.py Normal file
View 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