mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 17:00:23 +00:00
248 lines
7.5 KiB
Python
248 lines
7.5 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
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 "" in md
|
|
assert "## Legend" 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 "" in md
|
|
assert "## Caption" 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 "" in md
|
|
|
|
|
|
def test_render_figure_markdown_with_image_when_cropped() -> 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,
|
|
"was_cropped": True,
|
|
})
|
|
|
|
assert "" in md
|
|
|
|
|
|
def test_render_figure_markdown_without_image_when_not_cropped() -> 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,
|
|
"was_cropped": False,
|
|
})
|
|
|
|
assert "![](" not in md
|
|
assert "" not in md
|
|
|
|
|
|
def test_legend_only_matched_figure_does_not_use_unmatched_assets(tmp_path: Path) -> None:
|
|
from paperforge.worker.ocr_objects import extract_and_write_objects
|
|
|
|
render_root = tmp_path / "render"
|
|
asset_root = tmp_path / "assets"
|
|
|
|
figure_inventory: dict[str, Any] = {
|
|
"matched_figures": [
|
|
{
|
|
"text": "Figure 1. A legend-only figure.",
|
|
"page": 3,
|
|
"confidence": 0.85,
|
|
"cluster_bbox": None,
|
|
"matched_assets": [],
|
|
}
|
|
],
|
|
"unmatched_assets": [
|
|
{"bbox": [100, 100, 200, 200], "page": 1},
|
|
{"bbox": [300, 300, 400, 400], "page": 1},
|
|
],
|
|
"rejected_legends": [],
|
|
"figure_legends": [],
|
|
"figure_assets": [],
|
|
"official_figure_count": 1,
|
|
"unresolved_clusters": [],
|
|
}
|
|
|
|
extract_and_write_objects(
|
|
pdf_path=None,
|
|
figure_inventory=figure_inventory,
|
|
table_inventory={"tables": [], "unmatched_assets": []},
|
|
asset_root=asset_root,
|
|
render_root=render_root,
|
|
)
|
|
|
|
render_files = sorted((render_root / "figures").glob("*.md"))
|
|
assert len(render_files) == 3, (
|
|
"One matched figure note + two orphan notes from unmatched_assets"
|
|
)
|
|
figure_note = render_root / "figures" / "figure_001.md"
|
|
assert figure_note.exists()
|
|
content = figure_note.read_text()
|
|
assert "![](" not in content, (
|
|
"Legend-only figure must not contain any image reference"
|
|
)
|
|
# Verify orphan notes still get image paths
|
|
orphan_contents = (render_root / "figures" / "orphan_001.md").read_text()
|
|
assert "" in orphan_contents
|
|
|
|
|
|
def test_stabilize_object_wikilink_uses_correct_relative_path() -> None:
|
|
from paperforge.worker.ocr_objects import render_figure_object_markdown
|
|
|
|
md = render_figure_object_markdown({
|
|
"figure_id": "figure_001",
|
|
"page": 5,
|
|
"caption": "Figure 1. Results.",
|
|
"image_relpath": "assets/figures/figure_001.jpg",
|
|
"confidence": 0.9,
|
|
})
|
|
|
|
assert "" in md
|
|
|
|
|
|
def test_unresolved_cluster_object_emission(tmp_path: Path) -> None:
|
|
from paperforge.worker.ocr_objects import extract_and_write_objects
|
|
|
|
render_root = tmp_path / "render"
|
|
asset_root = tmp_path / "assets"
|
|
|
|
figure_inventory: dict[str, Any] = {
|
|
"matched_figures": [],
|
|
"unmatched_assets": [],
|
|
"rejected_legends": [],
|
|
"figure_legends": [],
|
|
"figure_assets": [],
|
|
"official_figure_count": 0,
|
|
"unresolved_clusters": [
|
|
{
|
|
"cluster_id": "cluster_001",
|
|
"page": 9,
|
|
"cluster_bbox": [363, 237, 1075, 1016],
|
|
"media_block_ids": ["p9_b2", "p9_b3", "p9_b4", "p9_b5", "p9_b6", "p9_b7"],
|
|
}
|
|
],
|
|
}
|
|
|
|
extract_and_write_objects(
|
|
pdf_path=None,
|
|
figure_inventory=figure_inventory,
|
|
table_inventory={"tables": [], "unmatched_assets": []},
|
|
asset_root=asset_root,
|
|
render_root=render_root,
|
|
)
|
|
|
|
render_files = sorted((render_root / "figures").glob("*.md"))
|
|
assert len(render_files) == 1, (
|
|
"Unresolved cluster should produce exactly one object note"
|
|
)
|
|
content = render_files[0].read_text()
|
|
assert "# Figure 4" not in content, (
|
|
"Cluster without a valid legend must not be titled with a figure number"
|
|
)
|
|
assert "cluster_001" in render_files[0].stem, (
|
|
"Object note should be identified by its cluster ID"
|
|
)
|
|
assert "cluster_001.jpg" in content, (
|
|
"Markdown must reference the whole-cluster crop image, not individual panel crops"
|
|
)
|
|
|
|
|
|
def test_crop_asset_uses_ocr_page_coordinates_when_dimensions_provided(tmp_path: Path) -> None:
|
|
import fitz
|
|
from PIL import Image
|
|
|
|
from paperforge.worker.ocr_objects import _crop_asset_from_pdf
|
|
|
|
pdf_path = tmp_path / "sample.pdf"
|
|
doc = fitz.open()
|
|
page = doc.new_page(width=300, height=400)
|
|
page.draw_rect(fitz.Rect(25, 25, 50, 50), color=(1, 0, 0), fill=(1, 0, 0))
|
|
doc.save(pdf_path)
|
|
doc.close()
|
|
|
|
dst = tmp_path / "crop.jpg"
|
|
ok = _crop_asset_from_pdf(
|
|
pdf_path,
|
|
1,
|
|
[50, 50, 100, 100],
|
|
dst,
|
|
page_width=600,
|
|
page_height=800,
|
|
page_cache_dir=tmp_path / "pages",
|
|
)
|
|
|
|
assert ok is True
|
|
with Image.open(dst) as img:
|
|
assert img.size == (50, 50)
|
|
|
|
|
|
def test_crop_asset_prefers_cached_page_image_when_available(tmp_path: Path) -> None:
|
|
from PIL import Image
|
|
|
|
from paperforge.worker.ocr_objects import _crop_asset_from_pdf
|
|
|
|
page_cache_dir = tmp_path / "pages"
|
|
page_cache_dir.mkdir()
|
|
page_image = page_cache_dir / "page_001.jpg"
|
|
Image.new("RGB", (600, 800), "white").save(page_image)
|
|
|
|
dst = tmp_path / "crop.jpg"
|
|
ok = _crop_asset_from_pdf(
|
|
tmp_path / "missing.pdf",
|
|
1,
|
|
[50, 50, 100, 100],
|
|
dst,
|
|
page_cache_dir=page_cache_dir,
|
|
)
|
|
|
|
assert ok is True
|
|
with Image.open(dst) as img:
|
|
assert img.size == (50, 50)
|