658 lines
22 KiB
Python
658 lines
22 KiB
Python
"""PDF 图片与表格提取 — 基于 pymupdf4llm layout analysis。
|
||
|
||
用 pymupdf4llm 的 layout analysis 检测 table / picture 区域，
|
||
再通过 caption 文字匹配确定 Figure/Table 编号，渲染为 JPEG。
|
||
|
||
相比旧方案（caption 正则 + pdfplumber/find_tables/文本块扫描三套策略）：
|
||
- layout analysis 直接给出区域 bbox，不存在相邻表格互相侵入的问题
|
||
- 无需手动调参（最大高度、间隙阈值等）
|
||
- 页面级 caption 匹配：每个 caption 只分配给最近的 box，避免上下相邻表格抢夺同一个 caption
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import re
|
||
from pathlib import Path
|
||
|
||
import pymupdf
|
||
import pymupdf4llm.helpers.document_layout as dl
|
||
|
||
from app.services.pdf_downloader import paper_dir
|
||
from app.utils import TMP_DIR
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ── Caption 正则 ───────────────────────────────────────────────────────
|
||
|
||
# 用于从 caption 文字中提取 Figure/Table 编号
|
||
_FIGURE_CAPTION_RE = re.compile(
|
||
r"^(?:Fig\.?|Figure)\s+(\d+)\s*(?:[:\.]\s*|\s+(?=(?-i:[A-Z])))",
|
||
re.IGNORECASE,
|
||
)
|
||
_TABLE_CAPTION_RE = re.compile(
|
||
r"^Table\s+(\d+)\s*(?:[:\.]\s*|\s+(?=(?-i:[A-Z])))",
|
||
re.IGNORECASE,
|
||
)
|
||
|
||
# caption 与 table/picture 的最大匹配距离（点）
|
||
_CAPTION_MATCH_DISTANCE = 100
|
||
# 截图区域的外边距
|
||
_REGION_PADDING = 5
|
||
# 3x 渲染，保证清晰度
|
||
_RENDER_ZOOM = 3
|
||
# 相邻 box 聚类间距（点）— 同一 figure/table 的碎片间距通常 < 15pt
|
||
_CLUSTER_GAP = 15
|
||
|
||
|
||
# ── Box 聚类 ─────────────────────────────────────────────────────────
|
||
|
||
|
||
class _BoxCluster:
|
||
"""合并后的布局区域（由一个或多个相邻 LayoutBox 组成）。
|
||
|
||
pymupdf4llm 有时将一个大图拆成多个小 picture box（如视频帧网格），
|
||
聚类后用整体 bbox 作为渲染区域。
|
||
"""
|
||
|
||
__slots__ = ("x0", "y0", "x1", "y1", "boxclass")
|
||
|
||
def __init__(self, boxes: list):
|
||
self.x0 = min(b.x0 for b in boxes)
|
||
self.y0 = min(b.y0 for b in boxes)
|
||
self.x1 = max(b.x1 for b in boxes)
|
||
self.y1 = max(b.y1 for b in boxes)
|
||
# table-fallback 归一化为 table（layout model 检测到表格但无法提取结构）
|
||
raw = boxes[0].boxclass
|
||
self.boxclass = "table" if raw == "table-fallback" else raw
|
||
|
||
|
||
def _cluster_boxes(boxes: list, gap: float = _CLUSTER_GAP) -> list[_BoxCluster]:
|
||
"""将相邻的同类型 box 合并为聚类。
|
||
|
||
用 union-find 将间距 ≤ gap 的同类型 box 归为一组，
|
||
每组生成一个 _BoxCluster（整体 bbox）。
|
||
"""
|
||
if not boxes:
|
||
return []
|
||
|
||
n = len(boxes)
|
||
parent = list(range(n))
|
||
|
||
def find(x: int) -> int:
|
||
while parent[x] != x:
|
||
parent[x] = parent[parent[x]]
|
||
x = parent[x]
|
||
return x
|
||
|
||
def union(a: int, b: int) -> None:
|
||
ra, rb = find(a), find(b)
|
||
if ra != rb:
|
||
parent[ra] = rb
|
||
|
||
for i in range(n):
|
||
bi = boxes[i]
|
||
for j in range(i + 1, n):
|
||
bj = boxes[j]
|
||
if bi.boxclass != bj.boxclass:
|
||
continue
|
||
h_gap = max(0.0, max(bi.x0, bj.x0) - min(bi.x1, bj.x1))
|
||
v_gap = max(0.0, max(bi.y0, bj.y0) - min(bi.y1, bj.y1))
|
||
h_overlap = bi.x1 > bj.x0 - gap and bj.x1 > bi.x0 - gap
|
||
v_overlap = bi.y1 > bj.y0 - gap and bj.y1 > bi.y0 - gap
|
||
if (h_gap <= gap and v_overlap) or (v_gap <= gap and h_overlap):
|
||
union(i, j)
|
||
|
||
groups: dict[int, list] = {}
|
||
for i in range(n):
|
||
groups.setdefault(find(i), []).append(boxes[i])
|
||
|
||
return [_BoxCluster(members) for members in groups.values()]
|
||
|
||
|
||
# ── 页面级 Caption 查找与匹配 ──────────────────────────────────────────
|
||
|
||
|
||
def _find_page_captions(page) -> list[dict]:
|
||
"""查找页面上所有 Figure/Table caption 文字块。"""
|
||
blocks = page.get_text("blocks")
|
||
captions = []
|
||
for b in blocks:
|
||
if len(b) < 5:
|
||
continue
|
||
bx0, by0, bx1, by1 = b[0], b[1], b[2], b[3]
|
||
text = str(b[4]).strip()
|
||
first_line = text.split("\n")[0].strip()
|
||
|
||
cap_type = None
|
||
m = _TABLE_CAPTION_RE.match(first_line)
|
||
if m:
|
||
cap_type = "table"
|
||
else:
|
||
m = _FIGURE_CAPTION_RE.match(first_line)
|
||
if m:
|
||
cap_type = "figure"
|
||
if m is None:
|
||
continue
|
||
|
||
captions.append(
|
||
{
|
||
"label": f"{'Table' if cap_type == 'table' else 'Figure'} {m.group(1)}",
|
||
"type": cap_type,
|
||
"caption_text": text,
|
||
"caption_y0": by0,
|
||
"caption_y1": by1,
|
||
"caption_x0": bx0,
|
||
"caption_x1": bx1,
|
||
}
|
||
)
|
||
return captions
|
||
|
||
|
||
def _vertical_distance(cap_y0, cap_y1, box_y0, box_y1) -> float | None:
|
||
"""计算 caption 到 box 的垂直距离。不邻接时返回 None。
|
||
|
||
三种情况：caption 完全在 box 上方、完全在下方、与 box 有垂直重叠。
|
||
重叠（含部分溢出）视为 distance=0，确保 caption 延伸到 box 边界外时不会丢失。
|
||
"""
|
||
# Caption 完全在 box 上方
|
||
if cap_y1 <= box_y0:
|
||
dist = box_y0 - cap_y1
|
||
return dist if dist <= _CAPTION_MATCH_DISTANCE else None
|
||
# Caption 完全在 box 下方
|
||
if cap_y0 >= box_y1:
|
||
dist = cap_y0 - box_y1
|
||
return dist if dist <= _CAPTION_MATCH_DISTANCE else None
|
||
# Caption 与 box 有垂直重叠（内部、部分溢出都算）→ 距离 0
|
||
return 0
|
||
|
||
|
||
def _same_column(cap: dict, box, page_width: float) -> bool:
|
||
"""判断 caption 和 box 是否在同一列。
|
||
|
||
双栏论文中左右栏间距有限，简单的水平重叠检查会跨列匹配。
|
||
策略：用中心 X 坐标判断各自在哪半边，只有同半边才算同列。
|
||
跨栏图表（caption 或 box 宽度 >65% 页宽）不受此限制。
|
||
"""
|
||
cap_w = cap["caption_x1"] - cap["caption_x0"]
|
||
box_w = box.x1 - box.x0
|
||
|
||
# 跨栏元素：宽度超过页面的 65%
|
||
if cap_w > page_width * 0.65 or box_w > page_width * 0.65:
|
||
return True
|
||
|
||
cap_cx = (cap["caption_x0"] + cap["caption_x1"]) / 2
|
||
box_cx = (box.x0 + box.x1) / 2
|
||
mid = page_width / 2
|
||
|
||
# 同在左半边或同在右半边
|
||
return (cap_cx < mid) == (box_cx < mid)
|
||
|
||
|
||
def _match_captions_to_boxes(
|
||
page_boxes: list, captions: list[dict], page_width: float
|
||
) -> list[tuple[list[int], list[dict]]]:
|
||
"""将 caption 分配给 box，允许一个 caption 匹配多个同类型 box。
|
||
|
||
典型场景：
|
||
- Figure 由左右两个 picture box 组成，caption 同时靠近两者
|
||
- Table 的视觉内容被 layout analysis 误分类为 picture，需要跨类型匹配
|
||
|
||
Returns:
|
||
[(box_indices, captions), ...] 每组是一个独立的渲染任务
|
||
"""
|
||
# 每个 caption 找到所有距离在阈值内的 box
|
||
# 优先匹配同类型；如果找不到，再匹配任意 table/picture box
|
||
cap_to_boxes: dict[int, list[tuple[int, float]]] = {}
|
||
|
||
for ci, cap in enumerate(captions):
|
||
same_type: list[tuple[int, float]] = []
|
||
any_type: list[tuple[int, float]] = []
|
||
expected = "table" if cap["type"] == "table" else "picture"
|
||
|
||
for bi, box in enumerate(page_boxes):
|
||
# 列感知：双栏论文中只匹配同栏的 box
|
||
if not _same_column(cap, box, page_width):
|
||
continue
|
||
# 水平重叠检查（同列内仍需有重叠）
|
||
if not (
|
||
cap["caption_x1"] > box.x0 - 5 and cap["caption_x0"] < box.x1 + 5
|
||
):
|
||
continue
|
||
dist = _vertical_distance(
|
||
cap["caption_y0"], cap["caption_y1"], box.y0, box.y1
|
||
)
|
||
if dist is None:
|
||
continue
|
||
entry = (bi, dist)
|
||
any_type.append(entry)
|
||
if box.boxclass == expected:
|
||
same_type.append(entry)
|
||
|
||
# 优先用同类型匹配；没有时回退到任意类型；都没有则跳过
|
||
if same_type:
|
||
cap_to_boxes[ci] = same_type
|
||
elif any_type:
|
||
cap_to_boxes[ci] = any_type
|
||
# else: 该 caption 无匹配 box，不加入 cap_to_boxes
|
||
|
||
# 每个 caption → 最近的 box（用于分组），但记录所有匹配的 box
|
||
cap_primary: dict[int, int] = {} # caption → primary box index
|
||
cap_all_boxes: dict[int, list[int]] = {} # caption → all matched box indices
|
||
for ci, matches in cap_to_boxes.items():
|
||
matches.sort(key=lambda x: x[1])
|
||
cap_primary[ci] = matches[0][0]
|
||
# 所有距离最近的同组 box（距离差 < 20pt 视为同一组）
|
||
best_dist = matches[0][1]
|
||
cap_all_boxes[ci] = [bi for bi, d in matches if d <= best_dist + 20]
|
||
|
||
# 按 primary box 分组
|
||
box_to_caps: dict[int, list[int]] = {}
|
||
for ci, bi in cap_primary.items():
|
||
box_to_caps.setdefault(bi, []).append(ci)
|
||
|
||
# 构建渲染组：每个 caption 独立成组（共享 box 但各自渲染）
|
||
# 同类型同 label 的 caption 会合并；不同类型则分开
|
||
used_captions: set[int] = set()
|
||
groups: list[tuple[list[int], list[dict]]] = []
|
||
|
||
for bi in sorted(box_to_caps.keys()):
|
||
cis = box_to_caps[bi]
|
||
for ci in cis:
|
||
if ci in used_captions:
|
||
continue
|
||
used_captions.add(ci)
|
||
|
||
all_box_indices = set(cap_all_boxes.get(ci, [bi]))
|
||
# 只合并同 label 的 caption（同 figure/table 的重复 caption）
|
||
merged_captions = [captions[ci]]
|
||
for other_bi in all_box_indices:
|
||
if other_bi in box_to_caps:
|
||
for other_ci in box_to_caps[other_bi]:
|
||
if other_ci not in used_captions:
|
||
other_cap = captions[other_ci]
|
||
if other_cap["label"] == captions[ci]["label"]:
|
||
used_captions.add(other_ci)
|
||
merged_captions.append(other_cap)
|
||
groups.append((sorted(all_box_indices), merged_captions))
|
||
|
||
return groups
|
||
|
||
|
||
# ── 单页处理 ─────────────────────────────────────────────────────────
|
||
|
||
|
||
def _render_and_save(
|
||
page,
|
||
clip: pymupdf.Rect,
|
||
images_dest: Path,
|
||
manifest: dict,
|
||
label: str,
|
||
cap_type: str,
|
||
caption_text: str,
|
||
page_num_1based: int,
|
||
arxiv_id: str,
|
||
) -> bool:
|
||
"""渲染页面区域并保存 JPEG，写入 manifest。成功返回 True。"""
|
||
mat = pymupdf.Matrix(_RENDER_ZOOM, _RENDER_ZOOM)
|
||
try:
|
||
pix = page.get_pixmap(matrix=mat, clip=clip)
|
||
except Exception:
|
||
logger.debug("Failed to render %s for %s", label, arxiv_id)
|
||
return False
|
||
|
||
filename = f"{label.replace(' ', '_').lower()}.jpg"
|
||
(images_dest / filename).write_bytes(pix.tobytes("jpeg"))
|
||
|
||
manifest[filename] = {
|
||
"page": page_num_1based,
|
||
"type": cap_type,
|
||
"label": label,
|
||
"caption_text": caption_text[:200] if caption_text else "",
|
||
"figures" if cap_type == "figure" else "tables": [label],
|
||
}
|
||
logger.debug(
|
||
"Rendered %s: page %d, region (%.0f,%.0f)-(%.0f,%.0f) → %s",
|
||
label,
|
||
page_num_1based,
|
||
clip.x0,
|
||
clip.y0,
|
||
clip.x1,
|
||
clip.y1,
|
||
filename,
|
||
)
|
||
return True
|
||
|
||
|
||
def _process_page(
|
||
doc,
|
||
page_idx: int,
|
||
page_layout,
|
||
images_dest: Path,
|
||
manifest: dict,
|
||
seen_labels: set,
|
||
arxiv_id: str,
|
||
) -> int:
|
||
"""处理单页：caption 匹配 + orphan 兜底，返回本页提取数量。"""
|
||
page = doc[page_idx]
|
||
page_width = page.rect.width
|
||
page_num = page_idx + 1
|
||
orphan_fig_counter = 0
|
||
orphan_tbl_counter = 0
|
||
|
||
# 收集本页的 table/picture box（跳过极小区域）
|
||
raw_boxes = []
|
||
for box in page_layout.boxes:
|
||
if box.boxclass not in ("table", "table-fallback", "picture"):
|
||
continue
|
||
if (box.x1 - box.x0) < 20 or (box.y1 - box.y0) < 20:
|
||
continue
|
||
raw_boxes.append(box)
|
||
|
||
if not raw_boxes:
|
||
return 0
|
||
|
||
# 聚类：将同一 figure/table 的碎片 box 合并
|
||
page_boxes = _cluster_boxes(raw_boxes)
|
||
|
||
# 页面级匹配：查找所有 caption，分配给 box
|
||
captions = _find_page_captions(page)
|
||
groups = _match_captions_to_boxes(page_boxes, captions, page_width)
|
||
|
||
# 只合并同 label 的 group（同一个 figure/table 的重复 caption）
|
||
# 不同 label 的 group 即使共享 box 也不合并（如 Figure 7 和 Figure 8），
|
||
# 渲染时用 caption 位置切割区域
|
||
_merged_groups: set[int] = set()
|
||
merged_groups: list[tuple[list[int], list[dict]]] = []
|
||
for gi, (box_indices, caps) in enumerate(groups):
|
||
if gi in _merged_groups:
|
||
continue
|
||
this_labels = {c["label"] for c in caps}
|
||
all_box_set = set(box_indices)
|
||
merge_targets = {gi}
|
||
for other_gi, (other_bi, other_caps) in enumerate(groups):
|
||
if other_gi <= gi or other_gi in _merged_groups:
|
||
continue
|
||
other_labels = {c["label"] for c in other_caps}
|
||
# 只在 label 有交集时合并（同一个 figure/table）
|
||
if this_labels & other_labels and all_box_set & set(other_bi):
|
||
merge_targets.add(other_gi)
|
||
all_box_set |= set(other_bi)
|
||
all_caps = []
|
||
for mgi in sorted(merge_targets):
|
||
_merged_groups.add(mgi)
|
||
all_caps.extend(groups[mgi][1])
|
||
merged_groups.append((sorted(all_box_set), all_caps))
|
||
groups = merged_groups
|
||
|
||
# ── 阶段 1：渲染有 caption 匹配的图/表 ──
|
||
matched_box_indices: set[int] = set()
|
||
extracted = 0
|
||
|
||
for box_indices, caps in groups:
|
||
matched_box_indices.update(box_indices)
|
||
|
||
# 去重同一 label，跳过已处理的
|
||
unique_caps = []
|
||
for cap in caps:
|
||
if cap["label"] not in seen_labels:
|
||
seen_labels.add(cap["label"])
|
||
unique_caps.append(cap)
|
||
if not unique_caps:
|
||
continue
|
||
|
||
# 合并所有关联 box 的 bbox
|
||
bx0 = min(page_boxes[i].x0 for i in box_indices)
|
||
by0 = min(page_boxes[i].y0 for i in box_indices)
|
||
bx1 = max(page_boxes[i].x1 for i in box_indices)
|
||
by1 = max(page_boxes[i].y1 for i in box_indices)
|
||
|
||
# 渲染区域：box + caption
|
||
all_cap_y0 = min(c["caption_y0"] for c in unique_caps)
|
||
all_cap_y1 = max(c["caption_y1"] for c in unique_caps)
|
||
all_cap_x0 = min(c["caption_x0"] for c in unique_caps)
|
||
all_cap_x1 = max(c["caption_x1"] for c in unique_caps)
|
||
|
||
top = max(0, min(by0, all_cap_y0) - _REGION_PADDING)
|
||
bottom = max(by1, all_cap_y1) + _REGION_PADDING
|
||
rx0 = max(0, min(bx0, all_cap_x0) - _REGION_PADDING)
|
||
rx1 = min(page_width, max(bx1, all_cap_x1) + _REGION_PADDING)
|
||
|
||
clip = pymupdf.Rect(rx0, top, rx1, bottom)
|
||
# 多个 caption 可能共享同一区域（如 subfigure），只需渲染一次
|
||
jpeg_bytes = None
|
||
for cap in unique_caps:
|
||
if jpeg_bytes is None:
|
||
if not _render_and_save(
|
||
page,
|
||
clip,
|
||
images_dest,
|
||
manifest,
|
||
cap["label"],
|
||
cap["type"],
|
||
cap["caption_text"],
|
||
page_num,
|
||
arxiv_id,
|
||
):
|
||
break
|
||
# 读取刚写入的 bytes 供后续同名 caption 复用
|
||
filename = f"{cap['label'].replace(' ', '_').lower()}.jpg"
|
||
jpeg_bytes = (images_dest / filename).read_bytes()
|
||
extracted += 1
|
||
else:
|
||
# 同区域的不同 caption（如 subfigure），复用图片
|
||
filename = f"{cap['label'].replace(' ', '_').lower()}.jpg"
|
||
(images_dest / filename).write_bytes(jpeg_bytes)
|
||
cap_preview = cap["caption_text"][:200]
|
||
manifest[filename] = {
|
||
"page": page_num,
|
||
"type": cap["type"],
|
||
"label": cap["label"],
|
||
"caption_text": cap_preview,
|
||
"figures" if cap["type"] == "figure" else "tables": [cap["label"]],
|
||
}
|
||
extracted += 1
|
||
|
||
# ── 阶段 2：渲染无 caption 匹配的图/表（orphan boxes） ──
|
||
orphan_indices = set(range(len(page_boxes))) - matched_box_indices
|
||
for bi in sorted(orphan_indices):
|
||
box = page_boxes[bi]
|
||
cap_type = "figure" if box.boxclass == "picture" else "table"
|
||
|
||
if cap_type == "figure":
|
||
orphan_fig_counter += 1
|
||
label = f"Figure (p{page_num}-{orphan_fig_counter})"
|
||
else:
|
||
orphan_tbl_counter += 1
|
||
label = f"Table (p{page_num}-{orphan_tbl_counter})"
|
||
|
||
if label in seen_labels:
|
||
continue
|
||
seen_labels.add(label)
|
||
|
||
clip = pymupdf.Rect(
|
||
max(0, box.x0 - _REGION_PADDING),
|
||
max(0, box.y0 - _REGION_PADDING),
|
||
min(page_width, box.x1 + _REGION_PADDING),
|
||
box.y1 + _REGION_PADDING,
|
||
)
|
||
if _render_and_save(
|
||
page,
|
||
clip,
|
||
images_dest,
|
||
manifest,
|
||
label,
|
||
cap_type,
|
||
"",
|
||
page_num,
|
||
arxiv_id,
|
||
):
|
||
extracted += 1
|
||
|
||
return extracted
|
||
|
||
|
||
# ── 核心提取 ───────────────────────────────────────────────────────────
|
||
|
||
|
||
def extract_images_from_pdf(arxiv_id: str, pdf_path: Path | None = None) -> int:
|
||
"""从 PDF 提取 Figure/Table 截图，生成 manifest。
|
||
|
||
用 pymupdf4llm layout analysis 检测 table/picture 区域，
|
||
再通过 caption 文字确定编号，渲染为 JPEG。
|
||
|
||
Args:
|
||
arxiv_id: 论文 ID
|
||
pdf_path: PDF 路径，默认 data/tmp/{arxiv_id}/paper.pdf
|
||
|
||
Returns:
|
||
提取的图片数量
|
||
"""
|
||
if pdf_path is None:
|
||
pdf_path = TMP_DIR / arxiv_id / "paper.pdf"
|
||
|
||
if not pdf_path.exists():
|
||
logger.warning("PDF not found for %s: %s", arxiv_id, pdf_path)
|
||
return 0
|
||
|
||
images_dest = paper_dir(arxiv_id) / "images"
|
||
images_dest.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 清理上次提取的旧图片
|
||
for old_file in images_dest.iterdir():
|
||
if old_file.suffix.lower() in (".png", ".jpg", ".jpeg"):
|
||
old_file.unlink()
|
||
if (images_dest / "manifest.json").exists():
|
||
(images_dest / "manifest.json").unlink()
|
||
|
||
doc = pymupdf.open(str(pdf_path))
|
||
|
||
# layout analysis
|
||
try:
|
||
parsed = dl.parse_document(
|
||
doc, filename=str(pdf_path), use_ocr=dl.OCRMode.NEVER
|
||
)
|
||
except Exception:
|
||
logger.warning(
|
||
"pymupdf4llm layout analysis failed for %s", arxiv_id, exc_info=True
|
||
)
|
||
doc.close()
|
||
return 0
|
||
|
||
extracted = 0
|
||
manifest: dict[str, dict] = {}
|
||
seen_labels: set[str] = set()
|
||
|
||
for page_idx, page_layout in enumerate(parsed.pages):
|
||
try:
|
||
extracted += _process_page(
|
||
doc,
|
||
page_idx,
|
||
page_layout,
|
||
images_dest=images_dest,
|
||
manifest=manifest,
|
||
seen_labels=seen_labels,
|
||
arxiv_id=arxiv_id,
|
||
)
|
||
except Exception:
|
||
logger.warning(
|
||
"Failed to process page %d for %s",
|
||
page_idx + 1,
|
||
arxiv_id,
|
||
exc_info=True,
|
||
)
|
||
continue
|
||
|
||
doc.close()
|
||
|
||
# 保存 manifest
|
||
manifest_path = images_dest / "manifest.json"
|
||
manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2))
|
||
|
||
if extracted > 0:
|
||
logger.info(
|
||
"Extracted %d figure/table screenshots from PDF for %s",
|
||
extracted,
|
||
arxiv_id,
|
||
)
|
||
|
||
return extracted
|
||
|
||
|
||
# ── 按 summary 过滤 ────────────────────────────────────────────────────
|
||
|
||
|
||
def filter_images_by_summary(arxiv_id: str, figures: list[dict]) -> int:
|
||
"""根据 summary 中的 figures 字段过滤提取的图片/表格。
|
||
|
||
用 manifest.json 中的 label 匹配，保留被 AI 总结引用的图片。
|
||
"""
|
||
if not figures:
|
||
return 0
|
||
|
||
images_dir = paper_dir(arxiv_id) / "images"
|
||
manifest_path = images_dir / "manifest.json"
|
||
|
||
if not images_dir.exists() or not manifest_path.exists():
|
||
return 0
|
||
|
||
all_files = [
|
||
f for f in images_dir.iterdir() if f.suffix.lower() in (".png", ".jpg", ".jpeg")
|
||
]
|
||
if not all_files:
|
||
return 0
|
||
|
||
manifest: dict = json.loads(manifest_path.read_text(encoding="utf-8"))
|
||
|
||
# 收集 summary 中引用的所有 Figure/Table ID（归一化）
|
||
referenced_ids: set[str] = set()
|
||
for fig in figures:
|
||
fig_id = fig.get("id", "")
|
||
m = re.match(r"(?:Fig\.?|Figure)\s*(\d+)", fig_id, re.IGNORECASE)
|
||
if m:
|
||
referenced_ids.add(f"Figure {m.group(1)}")
|
||
m2 = re.match(r"Table\s*(\d+)", fig_id, re.IGNORECASE)
|
||
if m2:
|
||
referenced_ids.add(f"Table {m2.group(1)}")
|
||
|
||
if not referenced_ids:
|
||
logger.warning("No valid figure/table IDs in summary for %s", arxiv_id)
|
||
return len(all_files)
|
||
|
||
# 根据 manifest 的 label 字段匹配
|
||
keep_filenames: set[str] = set()
|
||
for filename, info in manifest.items():
|
||
label = info.get("label", "")
|
||
if label in referenced_ids:
|
||
keep_filenames.add(filename)
|
||
continue
|
||
for ref in info.get("figures", []) + info.get("tables", []):
|
||
if ref in referenced_ids:
|
||
keep_filenames.add(filename)
|
||
break
|
||
|
||
if not keep_filenames:
|
||
logger.warning(
|
||
"No manifest matches for %s (refs=%s), keeping all",
|
||
arxiv_id,
|
||
referenced_ids,
|
||
)
|
||
return len(all_files)
|
||
|
||
removed = 0
|
||
for f in all_files:
|
||
if f.name not in keep_filenames:
|
||
f.unlink()
|
||
removed += 1
|
||
|
||
kept = len(all_files) - removed
|
||
logger.info(
|
||
"Filtered images for %s: kept %d, removed %d (refs=%s)",
|
||
arxiv_id,
|
||
kept,
|
||
removed,
|
||
referenced_ids,
|
||
)
|
||
return kept
|