from __future__ import annotations import contextlib import re from collections.abc import Callable from pathlib import Path from typing import Any from paperforge.worker.ocr_math import normalize_ocr_math_text def render_figure_object_markdown(figure: dict[str, Any]) -> str: caption = figure.get("caption", "") image_relpath = figure.get("image_relpath", "") # Extract figure number for the title figure_id = figure.get("figure_id", "") if figure_id.startswith("unresolved_cluster_") or figure_id.startswith("cluster_"): label = "Unresolved Figure Candidate" elif figure_id and not figure_id.startswith("orphan_"): m = re.search(r"\d+", figure_id) num = str(int(m.group())) if m else figure_id label = f"Figure {num}" else: label = "Orphan Media" was_cropped = figure.get("was_cropped", True) parts = [f"# {label}", ""] if was_cropped and image_relpath: parts.append(f"![](../../{image_relpath})") parts.append("") if caption: parts.append("## Legend") parts.append(normalize_ocr_math_text(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: caption = table.get("caption", "") image_relpath = table.get("image_relpath", "") formal_num = table.get("formal_table_number") if formal_num is not None: label = f"Table {formal_num}" else: table_id_raw = table.get("table_id", "unknown") m = re.search(r"\d+", table_id_raw) table_id = str(int(m.group())) if m else table_id_raw label = f"Table {table_id}" parts = [f"# {label}", "", f"![](../../{image_relpath})", ""] if caption: parts.append("## Caption") parts.append(normalize_ocr_math_text(caption)) note_texts = [normalize_ocr_math_text(t) for t in table.get("note_texts", []) if t] if note_texts: parts.append("") parts.append("## Notes") parts.extend(note_texts) 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 _find_cached_page_image(page_cache_dir: Path | None, page_num: int) -> Path | None: if page_cache_dir is None: return None for suffix in (".jpg", ".png"): candidate = page_cache_dir / f"page_{page_num:03d}{suffix}" if candidate.exists(): return candidate return None def _crop_asset_from_pdf( pdf_path: Path | None, page_num: int, bbox: list[float], dst: Path, *, page_width: int = 0, page_height: int = 0, page_cache_dir: Path | None = None, pdf_doc: Any | None = None, pdf_doc_provider: Callable[[], Any | None] | None = None, rotation_deg: int = 0, ) -> bool: if dst.exists(): with contextlib.suppress(Exception): dst.unlink() cached_page_image = _find_cached_page_image(page_cache_dir, page_num) if cached_page_image is not None and not rotation_deg: try: from paperforge.worker.ocr import crop_block_asset except ImportError: return False ok = crop_block_asset(cached_page_image, [int(v) for v in bbox], dst) return ok created_doc = None doc = pdf_doc if doc is None and pdf_doc_provider is not None: doc = pdf_doc_provider() if doc is None: return False if doc is None: if pdf_path is None or not pdf_path.exists(): return False try: import fitz except ImportError: return False created_doc = fitz.open(str(pdf_path)) doc = created_doc try: if page_width > 0 and page_height > 0 and page_cache_dir is not None and not rotation_deg: try: from paperforge.worker.ocr import crop_block_asset, render_pdf_page_cached except ImportError: return False try: page_image_path = page_cache_dir / f"page_{page_num:03d}.jpg" rendered = render_pdf_page_cached( doc, page_num, target_width=page_width, target_height=page_height, destination=page_image_path, ) if not rendered: return False ok = crop_block_asset(rendered, [int(v) for v in bbox], dst) return ok except Exception: return False try: import fitz except ImportError: return False try: page = doc[page_num - 1] # Clip is in PDF user space — bbox is in OCR coordinates (scaled). # Convert to PDF space using page_width/height vs PDF page rect ratio. pdf_rect = page.rect sx = max(1.0, page_width / pdf_rect.width) if page_width > 0 else 2.0 sy = max(1.0, page_height / pdf_rect.height) if page_height > 0 else 2.0 rect = fitz.Rect(bbox[0] / sx, bbox[1] / sy, bbox[2] / sx, bbox[3] / sy) # High-resolution zoom for crisp vector rendering (4x vs previous 2x) zoom = 4.0 mat = fitz.Matrix(zoom, zoom) pix = page.get_pixmap(matrix=mat, clip=rect) dst.parent.mkdir(parents=True, exist_ok=True) if rotation_deg: # PyMuPDF Pixmap has no .rotate() — render at 4x then rotate via PIL # on the high-res pixmap data. Single pass, no double JPEG compression. from PIL import Image import io img = Image.open(io.BytesIO(pix.tobytes("png"))) img = img.rotate(rotation_deg, expand=True, resample=Image.Resampling.LANCZOS) img.save(str(dst)) else: pix.save(str(dst)) return True except Exception: return False finally: if created_doc is not None: doc.close() 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, *, page_dimensions_by_page: dict[int, tuple[int, int]] | None = None, ) -> None: """Extract figure/table asset crops from PDF and write object markdown.""" # Validation-first figure matching may retain held figures in inventory, # but object emission remains limited to matched figures and unresolved # media clusters until figure evidence is sufficient. 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" page_cache_dir = asset_root.parent / "pages" for d in ( figures_asset_dir, tables_asset_dir, orphans_asset_dir, figures_render_dir, tables_render_dir, page_cache_dir, ): d.mkdir(parents=True, exist_ok=True) for stale_dir, pattern in ( (figures_asset_dir, "*.jpg"), (tables_asset_dir, "*.jpg"), (orphans_asset_dir, "*.jpg"), (figures_render_dir, "*.md"), (tables_render_dir, "*.md"), ): for stale in stale_dir.glob(pattern): with contextlib.suppress(Exception): stale.unlink() if page_dimensions_by_page is None: page_dimensions_by_page = {} def _page_dims(page_num: int) -> tuple[int, int]: return page_dimensions_by_page.get(page_num, (0, 0)) shared_pdf_doc: Any | None = None shared_pdf_open_attempted = False def _get_shared_pdf_doc() -> Any | None: nonlocal shared_pdf_doc, shared_pdf_open_attempted if shared_pdf_doc is not None: return shared_pdf_doc if shared_pdf_open_attempted: return None if pdf_path is None or not pdf_path.exists(): shared_pdf_open_attempted = True return None shared_pdf_open_attempted = True try: import fitz shared_pdf_doc = fitz.open(str(pdf_path)) except Exception: return None return shared_pdf_doc try: # Process matched figures for i, match in enumerate(figure_inventory.get("matched_figures", [])): fig_id = match.get("figure_id", f"figure_{i + 1:03d}") caption_text = match.get("text", "") page = match.get("page", 0) page_width, page_height = _page_dims(page) asset_path_rel = f"assets/figures/{fig_id}.jpg" asset_path_abs = figures_asset_dir / f"{fig_id}.jpg" rotation_deg = int(match.get("rotation_correction_deg", 0) or 0) was_cropped = False cluster_bbox = match.get("cluster_bbox") if cluster_bbox and all(v > 0 for v in cluster_bbox): was_cropped = _crop_asset_from_pdf( pdf_path, page, cluster_bbox, asset_path_abs, page_width=page_width, page_height=page_height, page_cache_dir=page_cache_dir, pdf_doc_provider=_get_shared_pdf_doc, rotation_deg=rotation_deg, ) if not was_cropped: 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) and _crop_asset_from_pdf( pdf_path, page, bbox, asset_path_abs, page_width=page_width, page_height=page_height, page_cache_dir=page_cache_dir, pdf_doc_provider=_get_shared_pdf_doc, rotation_deg=rotation_deg, ) ): 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), "was_cropped": was_cropped, } ) _write_object_markdown(md, figures_render_dir / f"{fig_id}.md") # Process unresolved figure clusters (multi-panel without reliable legend) for i, cluster in enumerate(figure_inventory.get("unresolved_clusters", [])): cluster_id = cluster.get("cluster_id") or f"unresolved_cluster_{i + 1:03d}" page = cluster.get("page", 0) bbox = cluster.get("cluster_bbox", [0, 0, 0, 0]) page_width, page_height = _page_dims(page) asset_path_rel = f"assets/figures/{cluster_id}.jpg" asset_path_abs = figures_asset_dir / f"{cluster_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, page_width=page_width, page_height=page_height, page_cache_dir=page_cache_dir, pdf_doc_provider=_get_shared_pdf_doc, ) md = render_figure_object_markdown( { "figure_id": cluster_id, "page": page, "caption": "", "image_relpath": asset_path_rel, "confidence": 0.45, } ) _write_object_markdown(md, figures_render_dir / f"{cluster_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) page_width, page_height = _page_dims(page) 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, page_width=page_width, page_height=page_height, page_cache_dir=page_cache_dir, pdf_doc_provider=_get_shared_pdf_doc, ) 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) page_width, page_height = _page_dims(page) 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 # Use render_bbox + render_rotation_deg for rotated tables _crop_bbox = table.get("render_bbox") or asset_bbox _rot_deg = table.get("render_rotation_deg", 0) or 0 if table.get("has_asset") and pdf_path and _crop_bbox and all(v > 0 for v in _crop_bbox): was_cropped = _crop_asset_from_pdf( pdf_path, page, _crop_bbox, asset_path_abs, page_width=page_width, page_height=page_height, page_cache_dir=page_cache_dir, pdf_doc_provider=_get_shared_pdf_doc, rotation_deg=_rot_deg, ) 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, "formal_table_number": table.get("formal_table_number") or table.get("table_number"), "note_texts": table.get("note_texts", []), "note_match_reason": table.get("note_match_reason", ""), } ) _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) page_width, page_height = _page_dims(page) 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, page_width=page_width, page_height=page_height, page_cache_dir=page_cache_dir, pdf_doc_provider=_get_shared_pdf_doc, ) 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") finally: if shared_pdf_doc is not None: with contextlib.suppress(Exception): shared_pdf_doc.close()