mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 06:50:53 +00:00
feat: emit OCR figure and table objects
This commit is contained in:
parent
cb79942044
commit
1de9cc281b
5 changed files with 337 additions and 17 deletions
|
|
@ -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(
|
||||
|
|
|
|||
229
paperforge/worker/ocr_objects.py
Normal file
229
paperforge/worker/ocr_objects.py
Normal 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")
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
49
tests/test_ocr_objects.py
Normal 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
|
||||
Loading…
Reference in a new issue