455 lines
14 KiB
Python
455 lines
14 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
from unittest.mock import MagicMock
|
||
|
||
import pymupdf
|
||
|
||
from app.services import pdf_image_extractor as mod
|
||
from app.services.layout_detector import LayoutBox
|
||
|
||
|
||
def _caption_block(bbox, text):
|
||
"""构造一个 page.get_text("dict") 风格的文本块。"""
|
||
return {
|
||
"type": 0,
|
||
"bbox": list(bbox),
|
||
"lines": [{"spans": [{"text": text}]}],
|
||
}
|
||
|
||
|
||
def test_process_page_pairs_caption_from_text(tmp_path):
|
||
"""caption 来自 PDF 文本流（figure 标题在内容下方），用其 ID 直接命名。"""
|
||
images_dest = tmp_path / "images"
|
||
images_dest.mkdir()
|
||
manifest: dict[str, dict] = {}
|
||
|
||
pix = MagicMock()
|
||
pix.tobytes.return_value = b"jpeg"
|
||
|
||
page = MagicMock()
|
||
page.rect.width = 600
|
||
page.rect.height = 800
|
||
page.get_pixmap.return_value = pix
|
||
page.get_text.return_value = {
|
||
"blocks": [
|
||
_caption_block((95, 310, 320, 325), "Figure 1: Overall architecture.")
|
||
]
|
||
}
|
||
|
||
doc = MagicMock()
|
||
doc.__getitem__.return_value = page
|
||
|
||
boxes = [LayoutBox(100, 100, 300, 300, "picture")]
|
||
|
||
extracted = mod._process_page(
|
||
doc,
|
||
0,
|
||
boxes,
|
||
images_dest=images_dest,
|
||
manifest=manifest,
|
||
seen_labels=set(),
|
||
arxiv_id="2401.00001",
|
||
)
|
||
|
||
assert extracted == 1
|
||
# caption 自带 ID → 直接命名 figure_1.jpg
|
||
info = manifest["figure_1.jpg"]
|
||
assert info["label"] == "Figure 1"
|
||
assert info["caption_text"] == "Figure 1: Overall architecture."
|
||
assert info["caption_source"] == "text"
|
||
|
||
|
||
def test_process_page_includes_caption_in_render(tmp_path):
|
||
"""渲染时把 caption 文本块区域合并进同一张截图。"""
|
||
images_dest = tmp_path / "images"
|
||
images_dest.mkdir()
|
||
manifest: dict[str, dict] = {}
|
||
|
||
pix = MagicMock()
|
||
pix.tobytes.return_value = b"jpeg"
|
||
|
||
page = MagicMock()
|
||
page.rect.width = 600
|
||
page.rect.height = 800
|
||
page.get_pixmap.return_value = pix
|
||
page.get_text.return_value = {
|
||
"blocks": [_caption_block((95, 310, 320, 325), "Figure 1: Caption text.")]
|
||
}
|
||
|
||
doc = MagicMock()
|
||
doc.__getitem__.return_value = page
|
||
|
||
boxes = [LayoutBox(100, 100, 300, 300, "picture")]
|
||
|
||
mod._process_page(
|
||
doc,
|
||
0,
|
||
boxes,
|
||
images_dest=images_dest,
|
||
manifest=manifest,
|
||
seen_labels=set(),
|
||
arxiv_id="2401.00001",
|
||
)
|
||
|
||
# 内容 [100,100,300,300] ∪ caption [95,310,320,325]，各方向加 _REGION_PADDING=5
|
||
# → Rect(90, 95, 325, 330)
|
||
clip = page.get_pixmap.call_args.kwargs["clip"]
|
||
assert clip == pymupdf.Rect(90, 95, 325, 330)
|
||
|
||
|
||
def test_process_page_table_caption_above(tmp_path):
|
||
"""table 标题惯例在内容上方，配对后命名 table_N.jpg。"""
|
||
images_dest = tmp_path / "images"
|
||
images_dest.mkdir()
|
||
manifest: dict[str, dict] = {}
|
||
|
||
pix = MagicMock()
|
||
pix.tobytes.return_value = b"jpeg"
|
||
|
||
page = MagicMock()
|
||
page.rect.width = 600
|
||
page.rect.height = 800
|
||
page.get_pixmap.return_value = pix
|
||
# caption 在内容上方 [80, 90, 320, 105]，内容表格 [80, 120, 320, 280]
|
||
page.get_text.return_value = {
|
||
"blocks": [_caption_block((80, 90, 320, 105), "Table 2 | Results summary.")]
|
||
}
|
||
|
||
doc = MagicMock()
|
||
doc.__getitem__.return_value = page
|
||
|
||
boxes = [LayoutBox(80, 120, 320, 280, "table")]
|
||
|
||
extracted = mod._process_page(
|
||
doc,
|
||
0,
|
||
boxes,
|
||
images_dest=images_dest,
|
||
manifest=manifest,
|
||
seen_labels=set(),
|
||
arxiv_id="2401.00001",
|
||
)
|
||
|
||
assert extracted == 1
|
||
info = manifest["table_2.jpg"]
|
||
assert info["label"] == "Table 2"
|
||
assert info["caption_source"] == "text"
|
||
|
||
|
||
def test_process_page_filters_uncaptioned(tmp_path):
|
||
"""没有 Figure/Table caption 配对的 box（Algorithm、无编号表等）被过滤，不输出。"""
|
||
images_dest = tmp_path / "images"
|
||
images_dest.mkdir()
|
||
manifest: dict[str, dict] = {}
|
||
|
||
pix = MagicMock()
|
||
pix.tobytes.return_value = b"jpeg"
|
||
|
||
page = MagicMock()
|
||
page.rect.width = 600
|
||
page.rect.height = 800
|
||
page.get_pixmap.return_value = pix
|
||
page.get_text.return_value = {"blocks": []} # 无任何 caption 文本块
|
||
|
||
doc = MagicMock()
|
||
doc.__getitem__.return_value = page
|
||
|
||
boxes = [LayoutBox(100, 100, 300, 300, "picture")]
|
||
|
||
extracted = mod._process_page(
|
||
doc,
|
||
0,
|
||
boxes,
|
||
images_dest=images_dest,
|
||
manifest=manifest,
|
||
seen_labels=set(),
|
||
arxiv_id="2401.00001",
|
||
)
|
||
|
||
assert extracted == 0
|
||
assert manifest == {}
|
||
|
||
|
||
def test_process_page_merges_subfigure_clusters(tmp_path):
|
||
"""复合图被切成多个稀疏子图框时，配到同一主 caption 后合并成一张截图。"""
|
||
images_dest = tmp_path / "images"
|
||
images_dest.mkdir()
|
||
manifest: dict[str, dict] = {}
|
||
|
||
pix = MagicMock()
|
||
pix.tobytes.return_value = b"jpeg"
|
||
|
||
page = MagicMock()
|
||
page.rect.width = 600
|
||
page.rect.height = 800
|
||
page.get_pixmap.return_value = pix
|
||
# 主 caption 覆盖整张复合图宽度，在内容下方
|
||
page.get_text.return_value = {
|
||
"blocks": [_caption_block((95, 410, 455, 425), "Figure 3: Composite figure.")]
|
||
}
|
||
|
||
doc = MagicMock()
|
||
doc.__getitem__.return_value = page
|
||
|
||
# 两个子图框左右排列，间距 50pt（> _CLUSTER_GAP=15，不被 _cluster_boxes 合并）
|
||
boxes = [
|
||
LayoutBox(100, 100, 250, 300, "picture"),
|
||
LayoutBox(300, 100, 450, 300, "picture"),
|
||
]
|
||
|
||
extracted = mod._process_page(
|
||
doc,
|
||
0,
|
||
boxes,
|
||
images_dest=images_dest,
|
||
manifest=manifest,
|
||
seen_labels=set(),
|
||
arxiv_id="2401.00003",
|
||
)
|
||
|
||
assert extracted == 1 # 两子图合并成一张
|
||
info = manifest["figure_3.jpg"]
|
||
assert info["label"] == "Figure 3"
|
||
assert info["subfigure_count"] == 2
|
||
# 合并 bbox = union 两子图 = [100,100,450,300]
|
||
assert info["box"] == [100.0, 100.0, 450.0, 300.0]
|
||
|
||
|
||
def test_process_page_two_tables_captions_below_not_merged(tmp_path):
|
||
"""同页相邻两张表、标题都在表下方时，各自配对自己的内容框，不互相吞并。
|
||
|
||
回归 2606.19926：Table 3/4 同页、两标题均在表下方。旧逻辑（table 标题
|
||
必在上的错向惩罚）会让 Table 3 吞掉 Table 4 的内容框、Table 4 缺图。
|
||
"""
|
||
images_dest = tmp_path / "images"
|
||
images_dest.mkdir()
|
||
manifest: dict[str, dict] = {}
|
||
|
||
pix = MagicMock()
|
||
pix.tobytes.return_value = b"jpeg"
|
||
|
||
page = MagicMock()
|
||
page.rect.width = 600
|
||
page.rect.height = 800
|
||
page.get_pixmap.return_value = pix
|
||
# 两个 table 标题都在各自内容下方（本文实际排版）
|
||
page.get_text.return_value = {
|
||
"blocks": [
|
||
_caption_block((80, 324, 460, 346), "Table 3 Pass@1 success rate."),
|
||
_caption_block((80, 503, 460, 547), "Table 4 Core gold-context metrics."),
|
||
]
|
||
}
|
||
|
||
doc = MagicMock()
|
||
doc.__getitem__.return_value = page
|
||
|
||
boxes = [
|
||
LayoutBox(80, 77, 460, 321, "table"), # Table 3 内容（上）
|
||
LayoutBox(80, 420, 460, 494, "table"), # Table 4 内容（下）
|
||
]
|
||
|
||
extracted = mod._process_page(
|
||
doc,
|
||
0,
|
||
boxes,
|
||
images_dest=images_dest,
|
||
manifest=manifest,
|
||
seen_labels=set(),
|
||
arxiv_id="2606.19926",
|
||
)
|
||
|
||
assert extracted == 2
|
||
assert "table_3.jpg" in manifest
|
||
assert "table_4.jpg" in manifest
|
||
# 各自只含自己的内容框，未被合并成一张大图
|
||
assert manifest["table_3.jpg"]["box"] == [80.0, 77.0, 460.0, 321.0]
|
||
assert manifest["table_4.jpg"]["box"] == [80.0, 420.0, 460.0, 494.0]
|
||
assert "subfigure_count" not in manifest["table_3.jpg"]
|
||
assert "subfigure_count" not in manifest["table_4.jpg"]
|
||
|
||
|
||
def test_caption_regex_rejects_multi_reference_body_text():
|
||
"""多图引用型正文（"Figure 14 and 15 show..."）不应被当作独立 caption。
|
||
|
||
回归 2606.24597：第 35 页正文 "Figure 14 and 15 show..." 被误判为 caption，
|
||
与真标题 "Figure 14:" 竞争 seen_labels 去重，导致 Figure 14 截成另一张图。
|
||
"""
|
||
re_ = mod._CAPTION_HEAD_RE
|
||
# 真标题：编号后接标题文字 / 冒号 / 子图编号，均应命中
|
||
assert re_.match("Figure 14: Interaction examples.")
|
||
assert re_.match("Figure 1 Context efficiency.")
|
||
assert re_.match("Table 1 Zero-shot results.")
|
||
assert re_.match("Figure 3.5 Sub-figure detail.")
|
||
# 多图引用正文：编号后紧跟 "and/to/, + 另一编号" → 拒绝
|
||
assert not re_.match("Figure 14 and 15 show representative examples.")
|
||
assert not re_.match("Figure 1 to 3 illustrate the pipeline.")
|
||
assert not re_.match("Table 2, 3 compare methods.")
|
||
|
||
|
||
def test_link_figures_fallback_excludes_assigned_images(tmp_path, monkeypatch):
|
||
"""序号兜底不复用已被精确匹配（策略 1）占用的图片。
|
||
|
||
回归 2606.19926：Table 4 在 manifest 缺图 → 兜底误取 table_1.jpg（已属 Table 1）。
|
||
"""
|
||
monkeypatch.setattr(mod, "PAPERS_DIR", tmp_path)
|
||
images_dir = tmp_path / "x" / "images"
|
||
images_dir.mkdir(parents=True)
|
||
# manifest 只命中 Table 1 → table_1.jpg；Table 4 无精确匹配，走兜底
|
||
(images_dir / "manifest.json").write_text(
|
||
json.dumps({"table_1.jpg": {"label": "Table 1"}})
|
||
)
|
||
|
||
figures = [{"id": "Table 1"}, {"id": "Table 4"}]
|
||
images = [
|
||
{"name": "table_1.jpg", "url": "/papers/x/images/table_1.jpg"},
|
||
{"name": "table_2.jpg", "url": "/papers/x/images/table_2.jpg"},
|
||
]
|
||
|
||
out = mod.link_figures_with_images(figures, images, "x")
|
||
by_id = {f["id"]: f for f in out}
|
||
|
||
assert by_id["Table 1"]["image_url"] == "/papers/x/images/table_1.jpg"
|
||
# 兜底只能拿到未被占用的 table_2.jpg，不能再复用 table_1.jpg
|
||
assert by_id["Table 4"]["image_url"] == "/papers/x/images/table_2.jpg"
|
||
|
||
|
||
def test_process_page_absorbs_far_subfigure_panel(tmp_path):
|
||
"""多面板图主 caption 只在最下方时，距 caption 过远的上方面板被吸收进来。
|
||
|
||
回归 2606.24597 Figure 4：(a) 子图距主标题 270pt（> _CAPTION_MATCH_DISTANCE），
|
||
原本被丢弃；(b) 子图距 46pt 正常配对。吸收后两面板合并成一张含 (a)(b)。
|
||
"""
|
||
images_dest = tmp_path / "images"
|
||
images_dest.mkdir()
|
||
manifest: dict[str, dict] = {}
|
||
|
||
pix = MagicMock()
|
||
pix.tobytes.return_value = b"jpeg"
|
||
page = MagicMock()
|
||
page.rect.width = 600
|
||
page.rect.height = 800
|
||
page.get_pixmap.return_value = pix
|
||
# 主 caption 在最下方 y=479，距上方 (a) 面板 270pt、距 (b) 面板 46pt
|
||
page.get_text.return_value = {
|
||
"blocks": [
|
||
_caption_block((80, 479, 520, 500), "Figure 4: Representative examples.")
|
||
]
|
||
}
|
||
doc = MagicMock()
|
||
doc.__getitem__.return_value = page
|
||
boxes = [
|
||
LayoutBox(80, 73, 520, 209, "picture"), # (a) 面板，距 caption 270pt
|
||
LayoutBox(80, 245, 520, 433, "picture"), # (b) 面板，距 caption 46pt
|
||
]
|
||
|
||
extracted = mod._process_page(
|
||
doc,
|
||
0,
|
||
boxes,
|
||
images_dest=images_dest,
|
||
manifest=manifest,
|
||
seen_labels=set(),
|
||
arxiv_id="2606.24597",
|
||
)
|
||
|
||
assert extracted == 1
|
||
info = manifest["figure_4.jpg"]
|
||
# 合并后含两面板：y 从 73 到 433
|
||
assert info["box"] == [80.0, 73.0, 520.0, 433.0]
|
||
assert info["subfigure_count"] == 2
|
||
|
||
|
||
def test_process_page_absorbs_chained_panels_multipass(tmp_path):
|
||
"""链式排列的多面板（上图→中图→下图→caption）需多趟吸收才能全部并入。
|
||
|
||
回归 2606.24597 Figure 15：3 个面板纵排、主 caption 只在最下方，中图被并入后
|
||
上图才能链上。单趟扫描会漏掉最上图。
|
||
"""
|
||
images_dest = tmp_path / "images"
|
||
images_dest.mkdir()
|
||
manifest: dict[str, dict] = {}
|
||
|
||
pix = MagicMock()
|
||
pix.tobytes.return_value = b"jpeg"
|
||
page = MagicMock()
|
||
page.rect.width = 600
|
||
page.rect.height = 800
|
||
page.get_pixmap.return_value = pix
|
||
# 主 caption 在最下方 y=710，仅下图（y505-669）在其 120pt 内
|
||
page.get_text.return_value = {
|
||
"blocks": [
|
||
_caption_block((80, 710, 520, 730), "Figure 15: Text-based domains.")
|
||
]
|
||
}
|
||
doc = MagicMock()
|
||
doc.__getitem__.return_value = page
|
||
boxes = [
|
||
LayoutBox(80, 73, 520, 217, "picture"), # 上图，距 caption 493pt
|
||
LayoutBox(80, 256, 520, 467, "picture"), # 中图，距 caption 243pt
|
||
LayoutBox(80, 505, 520, 669, "picture"), # 下图，距 caption 41pt（配得上）
|
||
]
|
||
|
||
extracted = mod._process_page(
|
||
doc,
|
||
0,
|
||
boxes,
|
||
images_dest=images_dest,
|
||
manifest=manifest,
|
||
seen_labels=set(),
|
||
arxiv_id="2606.24597",
|
||
)
|
||
|
||
assert extracted == 1
|
||
info = manifest["figure_15.jpg"]
|
||
# 三面板全并入：y 从 73 到 669
|
||
assert info["box"] == [80.0, 73.0, 520.0, 669.0]
|
||
assert info["subfigure_count"] == 3
|
||
|
||
|
||
def test_absorb_stragglers_respects_caption_boundary(tmp_path):
|
||
"""两图之间夹着（异类）标题时，漏配的上图不并入下图（吸收护栏 2）。
|
||
|
||
上图距下图标题 160pt（>120）漏配；下图正常配对。两图间夹一个 Table 标题
|
||
（异类，上图不会配它）→ 护栏判定有 caption 边界，不吸收上图。
|
||
"""
|
||
images_dest = tmp_path / "images"
|
||
images_dest.mkdir()
|
||
manifest: dict[str, dict] = {}
|
||
|
||
pix = MagicMock()
|
||
pix.tobytes.return_value = b"jpeg"
|
||
page = MagicMock()
|
||
page.rect.width = 600
|
||
page.rect.height = 800
|
||
page.get_pixmap.return_value = pix
|
||
page.get_text.return_value = {
|
||
"blocks": [
|
||
_caption_block(
|
||
(80, 160, 460, 175), "Table 9 Some table."
|
||
), # 夹中间的表标题
|
||
_caption_block((80, 310, 520, 330), "Figure 5 Some figure."), # 下图标题
|
||
]
|
||
}
|
||
doc = MagicMock()
|
||
doc.__getitem__.return_value = page
|
||
boxes = [
|
||
LayoutBox(80, 73, 520, 150, "picture"), # 上图，距 Figure 5 标题 160pt → 漏配
|
||
LayoutBox(80, 185, 520, 300, "picture"), # 下图，距 Figure 5 标题 10pt → 配上
|
||
]
|
||
|
||
extracted = mod._process_page(
|
||
doc,
|
||
0,
|
||
boxes,
|
||
images_dest=images_dest,
|
||
manifest=manifest,
|
||
seen_labels=set(),
|
||
arxiv_id="2606.24597",
|
||
)
|
||
|
||
# 上图未被并入下图 → 仅下图输出，box 不含上图、无 subfigure_count
|
||
assert extracted == 1
|
||
assert manifest["figure_5.jpg"]["box"] == [80.0, 185.0, 520.0, 300.0]
|
||
assert "subfigure_count" not in manifest["figure_5.jpg"]
|