feat: emit OCR figure and table objects

This commit is contained in:
Research Assistant 2026-06-05 00:34:22 +08:00
parent cb79942044
commit 1de9cc281b
5 changed files with 337 additions and 17 deletions

View file

@ -1779,6 +1779,20 @@ def postprocess_ocr_result(vault: Path, key: str, all_results: list[dict]) -> tu
table_inventory,
)
# --- Phase 2: object artifacts ---
from paperforge.worker.ocr_objects import extract_and_write_objects
ocr_asset_root = ocr_root / "assets"
ocr_render_root = ocr_root / "render"
extract_and_write_objects(
pdf_path=source_pdf_path,
figure_inventory=figure_inventory,
table_inventory=table_inventory,
asset_root=ocr_asset_root,
render_root=ocr_render_root,
)
# Update meta.json with version payloads
ocr_model = meta.get("ocr_model", meta.get("ocr_provider", "PaddleOCR"))
version_payload = build_version_payload(

View file

@ -0,0 +1,229 @@
from __future__ import annotations
from pathlib import Path
from typing import Any
def render_figure_object_markdown(figure: dict[str, Any]) -> str:
"""Render object markdown for a figure entry.
Produces Obsidian-flavored markdown with a wikilink embed to the
cropped asset image and the caption text.
"""
caption = figure.get("caption", "")
image_relpath = figure.get("image_relpath", "")
label = ""
if caption:
label = caption.split(".")[0].strip()
if not label:
figure_id = figure.get("figure_id", "")
if figure_id.startswith("orphan_"):
label = "Orphan Media"
else:
label = f"Figure {figure_id}"
parts = [f"# {label}", "", f"![[../{image_relpath}]]", ""]
if caption:
parts.append(caption)
if figure.get("page"):
parts.append("")
parts.append(f"*Page {figure['page']}*")
parts.append("")
parts.append("---")
return "\n".join(parts)
def render_table_object_markdown(table: dict[str, Any]) -> str:
"""Render object markdown for a table entry."""
caption = table.get("caption", "")
image_relpath = table.get("image_relpath", "")
label = ""
if caption:
label = caption.split(".")[0].strip()
if not label:
label = f"Table {table.get('table_id', 'unknown')}"
parts = [f"# {label}", "", f"![[../{image_relpath}]]", ""]
if caption:
parts.append(caption)
if table.get("page"):
parts.append("")
parts.append(f"*Page {table['page']}*")
parts.append("")
parts.append("---")
return "\n".join(parts)
def _write_object_markdown(md: str, dst: Path) -> None:
dst.parent.mkdir(parents=True, exist_ok=True)
dst.write_text(md.strip() + "\n", encoding="utf-8")
def _crop_asset_from_pdf(
pdf_path: Path,
page_num: int,
bbox: list[float],
dst: Path,
) -> bool:
"""Crop a region from a PDF page and save as JPEG.
Args:
pdf_path: Path to the PDF file.
page_num: 1-based page number.
bbox: [x1, y1, x2, y2] bounding box in PDF coordinates.
dst: Output path for the JPEG file.
Returns:
True if crop succeeded, False otherwise.
"""
try:
import fitz
except ImportError:
return False
if not pdf_path.exists():
return False
try:
doc = fitz.open(str(pdf_path))
page = doc[page_num - 1]
rect = fitz.Rect(*bbox)
pix = page.get_pixmap(matrix=fitz.Matrix(2, 2), clip=rect)
dst.parent.mkdir(parents=True, exist_ok=True)
pix.save(str(dst))
doc.close()
return True
except Exception:
return False
def extract_and_write_objects(
pdf_path: Path | None,
figure_inventory: dict[str, Any],
table_inventory: dict[str, Any],
asset_root: Path,
render_root: Path,
) -> None:
"""Extract figure/table asset crops from PDF and write object markdown.
Args:
pdf_path: Path to the source PDF (may be None).
figure_inventory: The figure inventory dict.
table_inventory: The table inventory dict.
asset_root: Root directory for assets (e.g., <ocr_root>/assets/).
render_root: Root directory for render objects (e.g., <ocr_root>/render/).
"""
figures_asset_dir = asset_root / "figures"
tables_asset_dir = asset_root / "tables"
orphans_asset_dir = asset_root / "orphans"
figures_render_dir = render_root / "figures"
tables_render_dir = render_root / "tables"
for d in (figures_asset_dir, tables_asset_dir, orphans_asset_dir,
figures_render_dir, tables_render_dir):
d.mkdir(parents=True, exist_ok=True)
# Process matched figures
for i, match in enumerate(figure_inventory.get("matched_figures", [])):
fig_id = f"figure_{i + 1:03d}"
caption_text = match.get("text", "")
page = match.get("page", 0)
asset_path_rel = f"assets/figures/{fig_id}.jpg"
asset_path_abs = figures_asset_dir / f"{fig_id}.jpg"
was_cropped = False
for asset_info in match.get("matched_assets", []):
bbox = asset_info.get("bbox", [0, 0, 0, 0])
if pdf_path and bbox and all(v > 0 for v in bbox):
if _crop_asset_from_pdf(pdf_path, page, bbox, asset_path_abs):
was_cropped = True
break
if not was_cropped:
for asset in figure_inventory.get("unmatched_assets", []):
bbox = asset.get("bbox", [0, 0, 0, 0])
asset_page = asset.get("page", 0)
if pdf_path and bbox and all(v > 0 for v in bbox):
if _crop_asset_from_pdf(pdf_path, asset_page, bbox, asset_path_abs):
was_cropped = True
break
md = render_figure_object_markdown({
"figure_id": fig_id,
"page": page,
"caption": caption_text,
"image_relpath": asset_path_rel,
"confidence": match.get("confidence", 0.5),
})
_write_object_markdown(md, figures_render_dir / f"{fig_id}.md")
# Process unmatched assets as orphans
orphan_count = 0
for asset in figure_inventory.get("unmatched_assets", []):
orphan_count += 1
orphan_id = f"orphan_{orphan_count:03d}"
bbox = asset.get("bbox", [0, 0, 0, 0])
page = asset.get("page", 0)
asset_path_rel = f"assets/orphans/{orphan_id}.jpg"
asset_path_abs = orphans_asset_dir / f"{orphan_id}.jpg"
was_cropped = False
if pdf_path and bbox and all(v > 0 for v in bbox):
was_cropped = _crop_asset_from_pdf(pdf_path, page, bbox, asset_path_abs)
md = render_figure_object_markdown({
"figure_id": orphan_id,
"page": page,
"caption": "",
"image_relpath": asset_path_rel,
"confidence": 0.3,
})
_write_object_markdown(md, figures_render_dir / f"{orphan_id}.md")
# Process tables
for i, table in enumerate(table_inventory.get("tables", [])):
tbl_id = f"table_{i + 1:03d}"
caption_text = table.get("caption_text", "")
page = table.get("page", 0)
asset_bbox = table.get("asset_bbox", [0, 0, 0, 0])
asset_path_rel = f"assets/tables/{tbl_id}.jpg"
asset_path_abs = tables_asset_dir / f"{tbl_id}.jpg"
was_cropped = False
if table.get("has_asset") and pdf_path and asset_bbox and all(v > 0 for v in asset_bbox):
was_cropped = _crop_asset_from_pdf(pdf_path, page, asset_bbox, asset_path_abs)
md = render_table_object_markdown({
"table_id": tbl_id,
"page": page,
"caption": caption_text,
"image_relpath": asset_path_rel,
"confidence": 0.85 if was_cropped else 0.4,
})
_write_object_markdown(md, tables_render_dir / f"{tbl_id}.md")
# Process unmatched table assets as orphans
for asset in table_inventory.get("unmatched_assets", []):
orphan_count += 1
orphan_id = f"orphan_{orphan_count:03d}"
bbox = asset.get("bbox", [0, 0, 0, 0])
page = asset.get("page", 0)
asset_path_rel = f"assets/orphans/{orphan_id}.jpg"
asset_path_abs = orphans_asset_dir / f"{orphan_id}.jpg"
if pdf_path and bbox and all(v > 0 for v in bbox):
_crop_asset_from_pdf(pdf_path, page, bbox, asset_path_abs)
md = render_figure_object_markdown({
"figure_id": orphan_id,
"page": page,
"caption": "",
"image_relpath": asset_path_rel,
"confidence": 0.3,
})
_write_object_markdown(md, figures_render_dir / f"{orphan_id}.md")

View file

@ -6,7 +6,6 @@ from typing import Any
from paperforge.core.io import write_json
_TABLE_PREFIX_PATTERN = re.compile(
r"^(?:Table|Supplementary\s+Table|Extended\s+Data\s+Table)\s+(\d+(?:\.\d+)?)",
flags=re.IGNORECASE,
@ -60,24 +59,22 @@ def build_table_inventory(structured_blocks: list[dict]) -> dict[str, Any]:
if matched_asset:
break
tables.append({
"caption_block_id": caption.get("block_id", ""),
"page": caption_page,
"caption_text": caption_text,
"table_number": table_num,
"asset_block_id": matched_asset.get("block_id", "") if matched_asset else "",
"asset_bbox": matched_asset.get("bbox", [0, 0, 0, 0]) if matched_asset else [],
"assistive_text": (matched_asset.get("text", "") or "") if matched_asset else "",
"truth_source": "image",
"has_asset": matched_asset is not None,
})
tables.append(
{
"caption_block_id": caption.get("block_id", ""),
"page": caption_page,
"caption_text": caption_text,
"table_number": table_num,
"asset_block_id": matched_asset.get("block_id", "") if matched_asset else "",
"asset_bbox": matched_asset.get("bbox", [0, 0, 0, 0]) if matched_asset else [],
"assistive_text": (matched_asset.get("text", "") or "") if matched_asset else "",
"truth_source": "image",
"has_asset": matched_asset is not None,
}
)
for caption in captions:
if not any(
t["caption_block_id"] == caption.get("block_id", "")
for t in tables
if t["has_asset"]
):
if not any(t["caption_block_id"] == caption.get("block_id", "") for t in tables if t["has_asset"]):
unmatched_captions.append(caption)
for i, asset in enumerate(assets):

View file

@ -379,3 +379,34 @@ def test_postprocess_writes_table_inventory(tmp_path: Path) -> None:
# Phase 2: table_inventory.json does not exist yet -- will fail
assert (ocr_dir / "structure" / "table_inventory.json").exists()
def test_postprocess_creates_object_directories(tmp_path: Path) -> None:
"""Verify postprocess creates object directories even with empty results."""
import json
from paperforge.worker.ocr import postprocess_ocr_result
vault = tmp_path / "vault"
vault.mkdir()
(vault / "paperforge.json").write_text(
json.dumps({"vault_config": {"system_dir": "System", "resources_dir": "Resources"}}),
encoding="utf-8",
)
ocr_root = vault / "System" / "PaperForge" / "ocr"
ocr_root.mkdir(parents=True)
ocr_dir = ocr_root / "OBJ001"
ocr_dir.mkdir()
(ocr_dir / "meta.json").write_text(
'{"zotero_key":"OBJ001","ocr_status":"done","ocr_model":"PaddleOCR"}',
encoding="utf-8",
)
_, _, _, _ = postprocess_ocr_result(vault, "OBJ001", [])
# Phase 2: object artifact directories are created during postprocess
assert (ocr_dir / "render" / "figures").exists()
assert (ocr_dir / "render" / "tables").exists()
assert (ocr_dir / "assets" / "figures").exists()
assert (ocr_dir / "assets" / "tables").exists()
assert (ocr_dir / "assets" / "orphans").exists()

49
tests/test_ocr_objects.py Normal file
View file

@ -0,0 +1,49 @@
from __future__ import annotations
def test_figure_object_markdown_links_image_and_legend() -> None:
from paperforge.worker.ocr_objects import render_figure_object_markdown
md = render_figure_object_markdown({
"figure_id": "figure_001",
"page": 4,
"caption": "Figure 1. Example.",
"image_relpath": "assets/figures/figure_001.jpg",
"confidence": 0.91,
})
assert "# Figure 1" in md
assert "![[../assets/figures/figure_001.jpg]]" in md
assert "Figure 1. Example." in md
def test_table_object_markdown_includes_image_and_caption() -> None:
from paperforge.worker.ocr_objects import render_table_object_markdown
md = render_table_object_markdown({
"table_id": "table_001",
"page": 5,
"caption": "Table 1. Results.",
"image_relpath": "assets/tables/table_001.jpg",
"confidence": 0.88,
})
assert "# Table 1" in md
assert "![[../assets/tables/table_001.jpg]]" in md
assert "Table 1. Results." in md
def test_orphan_object_markdown() -> None:
from paperforge.worker.ocr_objects import render_figure_object_markdown
md = render_figure_object_markdown({
"figure_id": "orphan_001",
"page": 6,
"caption": "",
"image_relpath": "assets/orphans/orphan_001.jpg",
"confidence": 0.3,
})
assert "# Orphan Media" in md
assert "![[../assets/orphans/orphan_001.jpg]]" in md