user2595/codex-py
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(F-09): rich parsing — formula + figure extraction

Browse Source 
- 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
Download Patch File Download Diff File Expand all files Collapse all files

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

311
codex/parsing/mathpix.py Normal file
View File

@@ -0,0 +1,311 @@
"""Formula extraction from PDF pages.
Two backends are supported:
1. **MathPix** (cloud): POST the PDF to ``api.mathpix.com/v3/pdf`` and poll
for the result JSON. Requires ``MATHPIX_APP_ID`` and ``MATHPIX_APP_KEY``
environment variables (or ``.env`` file).
2. **pix2tex** (local, MIT licence): crop each detected math bounding-box at
3× resolution and run the ``LatexOCR`` model locally. Used when MathPix
credentials are absent. Model (~116 MB) is downloaded on first call.
The public entry point :func:`extract_formulas` routes to whichever backend
is available.
"""
from __future__ import annotations
import logging
import unicodedata
from pathlib import Path
from typing import Any
from codex.config import get_settings
from codex.models import FormulaChunk
logger = logging.getLogger(__name__)
# Unicode categories that indicate mathematical symbols.
_MATH_CATEGORIES: frozenset[str] = frozenset({"Sm", "So"})
# Bbox size thresholds — avoids cropping tiny index fragments.
_MIN_HEIGHT_PT = 15.0
_MIN_WIDTH_PT = 40.0
# Minimum number of math-category characters in a text block before we
# consider it a formula candidate.
_MIN_MATH_SYMBOLS = 5
# Minimum ratio of math/total chars (as float 0–1).
_MIN_MATH_RATIO = 0.15
# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------
def _is_math_char(ch: str) -> bool:
"""Return True if *ch* belongs to a mathematical Unicode category."""
return unicodedata.category(ch) in _MATH_CATEGORIES
def _math_ratio(text: str) -> float:
"""Return the fraction of characters in *text* that are math chars."""
if not text:
return 0.0
math_count = sum(1 for ch in text if _is_math_char(ch))
return math_count / len(text)
def _math_symbol_count(text: str) -> int:
"""Return the absolute count of math-category characters in *text*."""
return sum(1 for ch in text if _is_math_char(ch))
# ---------------------------------------------------------------------------
# pix2tex singleton
# ---------------------------------------------------------------------------
_pix2tex_model: Any = None
def _get_pix2tex_model() -> Any:
"""Lazily load and return the pix2tex LatexOCR singleton."""
global _pix2tex_model # noqa: PLW0603
if _pix2tex_model is None:
from pix2tex.cli import LatexOCR
_pix2tex_model = LatexOCR()
return _pix2tex_model
# ---------------------------------------------------------------------------
# pix2tex backend
# ---------------------------------------------------------------------------
def _extract_formulas_pix2tex(pdf_path: str) -> list[FormulaChunk]:
"""Extract LaTeX formulas from *pdf_path* using pix2tex (local OCR).
Algorithm
---------
1. Open the PDF with PyMuPDF (fitz).
2. For each page, gather text blocks; keep only blocks whose
math-symbol ratio exceeds the threshold and whose bounding box
is larger than the minimum size.
3. Crop the region at 3× resolution (better OCR quality).
4. Run pix2tex ``LatexOCR`` on the cropped image.
5. Wrap the result in a :class:`~codex.models.FormulaChunk`.
"""
import fitz
from PIL import Image
results: list[FormulaChunk] = []
paper_id = Path(pdf_path).stem
model = _get_pix2tex_model()
try:
doc: Any = fitz.open(pdf_path)
except Exception:
logger.exception("fitz.open failed for %s", pdf_path)
return results
scale = fitz.Matrix(3, 3)
for page_num in range(len(doc)):
page = doc[page_num]
blocks: list[Any] = page.get_text("blocks")
for block in blocks:
# block = (x0, y0, x1, y1, text, block_no, block_type)
if len(block) < 5:
continue
x0, y0, x1, y1 = block[0], block[1], block[2], block[3]
text: str = block[4]
# Size filter
width_pt = x1 - x0
height_pt = y1 - y0
if width_pt < _MIN_WIDTH_PT or height_pt < _MIN_HEIGHT_PT:
continue
# Math content filter
if _math_symbol_count(text) < _MIN_MATH_SYMBOLS or _math_ratio(text) < _MIN_MATH_RATIO:
continue
# Crop at 3× for OCR quality
clip_rect = fitz.Rect(x0, y0, x1, y1)
crop_pix = page.get_pixmap(matrix=scale, clip=clip_rect)
img = Image.frombytes(
"RGB",
(crop_pix.width, crop_pix.height),
crop_pix.samples,
)
try:
latex: str = model(img)
except Exception:
logger.debug("pix2tex failed on page %d block, skipping", page_num + 1)
continue
if not latex or not latex.strip():
continue
context = text[:200].replace("\n", " ")
results.append(
FormulaChunk(
paper_id=paper_id,
page=page_num + 1,
raw_latex=latex.strip(),
context=context,
)
)
return results
# ---------------------------------------------------------------------------
# MathPix backend
# ---------------------------------------------------------------------------
def _extract_formulas_mathpix(
pdf_path: str,
app_id: str,
app_key: str,
) -> list[FormulaChunk]:
"""Extract LaTeX formulas using the MathPix v3 PDF API.
The function uploads the PDF, polls for results and parses the
returned ``latex_simplified`` list into :class:`~codex.models.FormulaChunk`
objects.
Retries with exponential back-off (via ``tenacity``) on transient HTTP
errors.
"""
import httpx
from tenacity import retry, stop_after_attempt, wait_exponential
paper_id = Path(pdf_path).stem
headers = {
"app_id": app_id,
"app_key": app_key,
}
@retry(
stop=stop_after_attempt(4),
wait=wait_exponential(multiplier=1, min=2, max=30),
reraise=True,
)
def _post_pdf() -> dict[str, Any]:
with open(pdf_path, "rb") as fh:
resp = httpx.post(
"https://api.mathpix.com/v3/pdf",
headers=headers,
files={"file": fh},
timeout=60.0,
)
resp.raise_for_status()
return resp.json() # type: ignore[no-any-return]
submit_response = _post_pdf()
pdf_id: str = submit_response.get("pdf_id", "")
if not pdf_id:
logger.warning("MathPix did not return a pdf_id — aborting")
return []
# Poll until complete
import time
for _ in range(60):
poll_resp = httpx.get(
f"https://api.mathpix.com/v3/pdf/{pdf_id}.mmd",
headers=headers,
timeout=30.0,
)
if poll_resp.status_code == 200:
break
time.sleep(2)
else:
logger.warning("MathPix polling timed out for pdf_id=%s", pdf_id)
return []
# Parse response — MathPix MMD contains \[ ... \] or $...$ blocks
mmd_text: str = poll_resp.text
results: list[FormulaChunk] = []
lines = mmd_text.splitlines()
in_block = False
block_lines: list[str] = []
page = 1
for line in lines:
if line.strip().startswith("\\["):
in_block = True
block_lines = [line]
elif in_block:
block_lines.append(line)
if line.strip().endswith("\\]"):
latex = "\n".join(block_lines).strip()
results.append(
FormulaChunk(
paper_id=paper_id,
page=page,
raw_latex=latex,
context="",
)
)
in_block = False
block_lines = []
# crude page tracking via MMD page markers
if line.strip().startswith("<!-- Page"):
page += 1
return results
# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------
def extract_formulas(
pdf_path: str,
mathpix_app_id: str | None = None,
mathpix_app_key: str | None = None,
) -> list[FormulaChunk]:
"""Extract mathematical formulas from *pdf_path*.
Routes to the MathPix cloud backend when ``mathpix_app_id`` and
``mathpix_app_key`` are both non-empty; otherwise falls back to
the local pix2tex model (if ``PIX2TEX_FALLBACK=true``, which is
the default).
Parameters
----------
pdf_path:
Absolute or relative path to the PDF file.
mathpix_app_id:
Override for ``MATHPIX_APP_ID`` setting (optional).
mathpix_app_key:
Override for ``MATHPIX_APP_KEY`` setting (optional).
Returns
-------
list[FormulaChunk]
Extracted formulas; empty list if extraction is disabled or fails.
"""
settings = get_settings()
app_id = mathpix_app_id or settings.mathpix_app_id or ""
app_key = mathpix_app_key or settings.mathpix_app_key or ""
if app_id and app_key:
logger.info("Using MathPix backend for %s", pdf_path)
return _extract_formulas_mathpix(pdf_path, app_id, app_key)
if settings.pix2tex_fallback:
logger.info("Using pix2tex fallback backend for %s", pdf_path)
return _extract_formulas_pix2tex(pdf_path)
logger.info("Formula extraction disabled (no creds, pix2tex_fallback=False)")
return []