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

Merged
user2595 merged 6 commits from feat/F-09-rich-parsing into main 2026-06-14 01:52:24 +00:00
3 changed files with 88 additions and 75 deletions
Showing only changes of commit cd42e9abc8 - Show all commits

View File

@@ -220,13 +220,18 @@ def ingest_paper(
if formulas or figures:
with get_conn() as conn2:
with conn2.cursor() as cur:
# Delete-before-insert keeps re-ingest idempotent (BIGSERIAL has no
# natural UNIQUE key, so ON CONFLICT DO NOTHING never fires).
cur.execute("DELETE FROM formulas WHERE paper_id = %s", (paper.id,))
cur.execute("DELETE FROM figures WHERE paper_id = %s", (paper.id,))
if formulas:
with conn2.cursor() as cur:
cur.executemany(
"""
INSERT INTO formulas (paper_id, page, raw_latex, context, eq_label)
VALUES (%s, %s, %s, %s, %s)
ON CONFLICT DO NOTHING
""",
[
(f.paper_id, f.page, f.raw_latex, f.context, f.eq_label)
@@ -241,7 +246,6 @@ def ingest_paper(
"""
INSERT INTO figures (paper_id, page, caption, image_path)
VALUES (%s, %s, %s, %s)
ON CONFLICT DO NOTHING
""",
[
(fig.paper_id, fig.page, fig.caption, fig.image_path)

View File

@@ -126,41 +126,44 @@ def extract_figures(
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")
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]
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]
# 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
# 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)
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,
results.append(
FigureChunk(
paper_id=paper_id,
page=page_num + 1,
image_path=str(img_path),
caption=caption,
)
)
)
finally:
doc.close()
return results

View File

@@ -112,54 +112,60 @@ def _extract_formulas_pix2tex(pdf_path: str) -> list[FormulaChunk]:
scale = fitz.Matrix(3, 3)
for page_num in range(len(doc)):
page = doc[page_num]
blocks: list[Any] = page.get_text("blocks")
try:
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]
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
# 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
# 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,
# 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,
)
)
finally:
doc.close()
return results