diff --git a/paperforge/worker/ocr.py b/paperforge/worker/ocr.py index e5e680d5..e483f7f0 100644 --- a/paperforge/worker/ocr.py +++ b/paperforge/worker/ocr.py @@ -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( diff --git a/paperforge/worker/ocr_objects.py b/paperforge/worker/ocr_objects.py new file mode 100644 index 00000000..532c83b3 --- /dev/null +++ b/paperforge/worker/ocr_objects.py @@ -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., /assets/). + render_root: Root directory for render objects (e.g., /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") diff --git a/paperforge/worker/ocr_tables.py b/paperforge/worker/ocr_tables.py index 3f44d05c..d610a78d 100644 --- a/paperforge/worker/ocr_tables.py +++ b/paperforge/worker/ocr_tables.py @@ -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): diff --git a/tests/test_ocr.py b/tests/test_ocr.py index 3205d478..d3f9b676 100644 --- a/tests/test_ocr.py +++ b/tests/test_ocr.py @@ -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() diff --git a/tests/test_ocr_objects.py b/tests/test_ocr_objects.py new file mode 100644 index 00000000..526f36e6 --- /dev/null +++ b/tests/test_ocr_objects.py @@ -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