refactor: improve PDF extraction, error handling, and proxy configuration
This commit is contained in:
@@ -23,13 +23,5 @@ class ValidationError(AppError):
|
|||||||
"""请求参数校验失败(400)。"""
|
"""请求参数校验失败(400)。"""
|
||||||
|
|
||||||
|
|
||||||
class ExternalAPIError(AppError):
|
|
||||||
"""外部 API 调用失败(502)。"""
|
|
||||||
|
|
||||||
|
|
||||||
class PdfProcessError(AppError):
|
|
||||||
"""PDF 处理错误(500)。"""
|
|
||||||
|
|
||||||
|
|
||||||
class ConflictError(AppError):
|
class ConflictError(AppError):
|
||||||
"""资源冲突(409)— 如锁冲突、并发任务冲突。"""
|
"""资源冲突(409)— 如锁冲突、并发任务冲突。"""
|
||||||
|
|||||||
-10
@@ -13,9 +13,7 @@ from app.config import settings
|
|||||||
from app.exceptions import (
|
from app.exceptions import (
|
||||||
AppError,
|
AppError,
|
||||||
ConflictError,
|
ConflictError,
|
||||||
ExternalAPIError,
|
|
||||||
NotFoundError,
|
NotFoundError,
|
||||||
PdfProcessError,
|
|
||||||
ValidationError,
|
ValidationError,
|
||||||
)
|
)
|
||||||
from app.database import engine, init_db
|
from app.database import engine, init_db
|
||||||
@@ -87,14 +85,6 @@ def create_app() -> FastAPI:
|
|||||||
async def _validation_handler(request, exc):
|
async def _validation_handler(request, exc):
|
||||||
return JSONResponse(status_code=400, content={"error": exc.message})
|
return JSONResponse(status_code=400, content={"error": exc.message})
|
||||||
|
|
||||||
@app.exception_handler(ExternalAPIError)
|
|
||||||
async def _external_api_handler(request, exc):
|
|
||||||
return JSONResponse(status_code=502, content={"error": exc.message})
|
|
||||||
|
|
||||||
@app.exception_handler(PdfProcessError)
|
|
||||||
async def _pdf_process_handler(request, exc):
|
|
||||||
return JSONResponse(status_code=500, content={"error": exc.message})
|
|
||||||
|
|
||||||
@app.exception_handler(ConflictError)
|
@app.exception_handler(ConflictError)
|
||||||
async def _conflict_handler(request, exc):
|
async def _conflict_handler(request, exc):
|
||||||
return JSONResponse(status_code=409, content={"error": exc.message})
|
return JSONResponse(status_code=409, content={"error": exc.message})
|
||||||
|
|||||||
@@ -41,7 +41,8 @@ async def call_claude(
|
|||||||
fix_errors: 上一轮验证错误列表(用于重试)
|
fix_errors: 上一轮验证错误列表(用于重试)
|
||||||
"""
|
"""
|
||||||
if session_id is None:
|
if session_id is None:
|
||||||
session_id = f"claude-summary-{uuid.uuid4().hex[:8]}"
|
# claude CLI 的 --session-id 要求合法 UUID;非 UUID 会被拒绝(Invalid session ID)
|
||||||
|
session_id = str(uuid.uuid4())
|
||||||
|
|
||||||
cmd = [settings.CLAUDE_BIN, "-p", "--output-format", "text"]
|
cmd = [settings.CLAUDE_BIN, "-p", "--output-format", "text"]
|
||||||
|
|
||||||
@@ -51,22 +52,22 @@ async def call_claude(
|
|||||||
else:
|
else:
|
||||||
cmd += ["--session-id", session_id]
|
cmd += ["--session-id", session_id]
|
||||||
|
|
||||||
cmd.append(prompt)
|
|
||||||
|
|
||||||
logger.info(
|
logger.info(
|
||||||
"Calling claude (session=%s, fix=%s)",
|
"Calling claude (session=%s, fix=%s)",
|
||||||
session_id,
|
session_id,
|
||||||
bool(fix_errors),
|
bool(fix_errors),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# prompt 走 stdin,避免长文本(论文 PDF 动辄数十 KB)超出系统 ARG_MAX
|
||||||
proc = await asyncio.create_subprocess_exec(
|
proc = await asyncio.create_subprocess_exec(
|
||||||
*cmd,
|
*cmd,
|
||||||
|
stdin=asyncio.subprocess.PIPE,
|
||||||
stdout=asyncio.subprocess.PIPE,
|
stdout=asyncio.subprocess.PIPE,
|
||||||
stderr=asyncio.subprocess.PIPE,
|
stderr=asyncio.subprocess.PIPE,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
stdout, stderr = await asyncio.wait_for(
|
stdout, stderr = await asyncio.wait_for(
|
||||||
proc.communicate(),
|
proc.communicate(input=prompt.encode("utf-8")),
|
||||||
timeout=settings.SUMMARY_TIMEOUT_SECONDS,
|
timeout=settings.SUMMARY_TIMEOUT_SECONDS,
|
||||||
)
|
)
|
||||||
except asyncio.TimeoutError:
|
except asyncio.TimeoutError:
|
||||||
@@ -77,8 +78,10 @@ async def call_claude(
|
|||||||
)
|
)
|
||||||
|
|
||||||
if proc.returncode != 0:
|
if proc.returncode != 0:
|
||||||
raise ClaudeProcessError(
|
# claude CLI 把 API 错误输出到 stdout(stderr 常为空),优先用有内容的
|
||||||
proc.returncode, stderr.decode("utf-8", errors="replace")
|
detail = stderr.decode("utf-8", errors="replace")
|
||||||
)
|
if not detail.strip():
|
||||||
|
detail = stdout.decode("utf-8", errors="replace")
|
||||||
|
raise ClaudeProcessError(proc.returncode, detail)
|
||||||
|
|
||||||
return stdout.decode("utf-8", errors="replace"), session_id
|
return stdout.decode("utf-8", errors="replace"), session_id
|
||||||
|
|||||||
@@ -140,7 +140,7 @@ def _get_embedding(text: str) -> list[float] | None:
|
|||||||
}
|
}
|
||||||
|
|
||||||
try:
|
try:
|
||||||
with make_http_client(sync=True) as client:
|
with make_http_client(sync=True, use_proxy=False) as client:
|
||||||
resp = client.post(url, json=payload, headers=headers)
|
resp = client.post(url, json=payload, headers=headers)
|
||||||
resp.raise_for_status()
|
resp.raise_for_status()
|
||||||
data = resp.json()
|
data = resp.json()
|
||||||
|
|||||||
@@ -36,12 +36,18 @@ _CLUSTER_GAP = 15
|
|||||||
_MIN_BOX_AREA = 2000
|
_MIN_BOX_AREA = 2000
|
||||||
# caption 文本块与 figure/table 内容块的最大垂直距离(单位: pt)
|
# caption 文本块与 figure/table 内容块的最大垂直距离(单位: pt)
|
||||||
_CAPTION_MATCH_DISTANCE = 120
|
_CAPTION_MATCH_DISTANCE = 120
|
||||||
# 方向不符(figure 标题在上 / table 标题在下)的配对惩罚分(仍允许,兜底异常排版)
|
# 游离碎片(配不到 caption)并入紧邻已配 cluster 的最大垂直间距(单位: pt)
|
||||||
_CAPTION_WRONG_SIDE_PENALTY = 300
|
# 容下多面板图的 "(a)/(b)" 子图标占位(实测 Figure 4 两面板间距 36pt)
|
||||||
|
_ABSORB_GAP = 60
|
||||||
# caption 开头标记:Figure 3 / Fig. 3 / Table C1 / Figure 3.5 等(大小写均可)
|
# caption 开头标记:Figure 3 / Fig. 3 / Table C1 / Figure 3.5 等(大小写均可)
|
||||||
# 编号 = 数字开头 或 字母+数字(附录 C1);行首匹配,规避正文 "see Table 3" 引用
|
# 编号 = 数字开头 或 字母+数字(附录 C1);行首匹配,规避正文 "see Table 3" 引用。
|
||||||
|
# 否定前瞻再排除多图引用型正文 —— "Figure 14 and 15 show..." / "Figure 1 to 3" /
|
||||||
|
# "Table 2, 3" 这类引用多张图表的句子不是独立标题,真标题从不引用多个编号。
|
||||||
_CAPTION_HEAD_RE = re.compile(
|
_CAPTION_HEAD_RE = re.compile(
|
||||||
r"^\s*(Figure|Fig\.?|Table)\b\.?\s+([0-9][0-9A-Za-z.]*|[A-Z]\d[0-9A-Za-z.]*)",
|
r"^\s*(Figure|Fig\.?|Table)\b\.?\s+([0-9][0-9A-Za-z.]*|[A-Z]\d[0-9A-Za-z.]*)"
|
||||||
|
r"(?![0-9A-Za-z.])" # 编号须是完整 token,防止 "Figure 14" 回溯成 "Figure 1" 逃逸
|
||||||
|
r"(?!\s*(?:and|to|through|vs\.?|&)\s+\d)" # "Figure 3 and 4" / "Figure 1 to 3"
|
||||||
|
r"(?!\s*,\s*\d)", # "Figure 3, 4"
|
||||||
re.IGNORECASE,
|
re.IGNORECASE,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -142,8 +148,8 @@ def _find_caption_blocks(page) -> list[_CaptionBlock]:
|
|||||||
"".join(span.get("text", "") for span in line.get("spans", []))
|
"".join(span.get("text", "") for span in line.get("spans", []))
|
||||||
for line in lines
|
for line in lines
|
||||||
]
|
]
|
||||||
first_line = next((t for t in line_texts if t.strip()), "")
|
joined = " ".join(t.strip() for t in line_texts if t.strip())
|
||||||
m = _CAPTION_HEAD_RE.match(first_line)
|
m = _CAPTION_HEAD_RE.match(joined)
|
||||||
if not m:
|
if not m:
|
||||||
continue
|
continue
|
||||||
kind_word, num = m.group(1), m.group(2)
|
kind_word, num = m.group(1), m.group(2)
|
||||||
@@ -151,7 +157,7 @@ def _find_caption_blocks(page) -> list[_CaptionBlock]:
|
|||||||
bbox = block.get("bbox")
|
bbox = block.get("bbox")
|
||||||
if not bbox or len(bbox) != 4:
|
if not bbox or len(bbox) != 4:
|
||||||
continue
|
continue
|
||||||
full_text = " ".join(t.strip() for t in line_texts if t.strip())
|
full_text = joined
|
||||||
results.append(
|
results.append(
|
||||||
_CaptionBlock(
|
_CaptionBlock(
|
||||||
id=f"{'Table' if is_table else 'Figure'} {num}",
|
id=f"{'Table' if is_table else 'Figure'} {num}",
|
||||||
@@ -166,18 +172,25 @@ def _find_caption_blocks(page) -> list[_CaptionBlock]:
|
|||||||
def _pair_caption_blocks(
|
def _pair_caption_blocks(
|
||||||
content_clusters: list[_BoxCluster],
|
content_clusters: list[_BoxCluster],
|
||||||
caption_blocks: list[_CaptionBlock],
|
caption_blocks: list[_CaptionBlock],
|
||||||
) -> dict[int, _CaptionBlock]:
|
) -> dict[int, list[int]]:
|
||||||
"""每个内容块配方向上最近的同类型标题块。
|
"""每个 caption 配对其垂直 span 内的所有同类型 cluster(支持复合图/子表)。
|
||||||
|
|
||||||
figure 标题惯例在下方、table 标题在上方;方向相符优先,不符加惩罚兜底
|
不预设标题在内容上方还是下方 —— figure 惯例标题在下、table 惯例标题在上,
|
||||||
(跨页 / 异常排版)。按 (距离+惩罚) 升序贪心匹配,每个内容块与标题块唯一配对。
|
但不少论文反向排版(table 标题在表下方);用方向作硬约束或加错向惩罚会把
|
||||||
|
"同页相邻两张表" 错并(标题居中的那张吞掉邻居)。改为上下两侧平等地按垂直
|
||||||
|
距离打分,每个 cluster 唯一归属最近的同类 caption,但一个 caption 可被多个
|
||||||
|
cluster 共享 —— 这样一张被 DocLayout 切成多个稀疏子框的复合图/复合表,能
|
||||||
|
整体配到它的主标题(而非只截其中一个子图/子表)。
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
caption_idx → [cluster_idx, ...],按 cluster 在页面上的位置排序,
|
||||||
|
保证合并/渲染顺序稳定。
|
||||||
"""
|
"""
|
||||||
candidates: list[tuple[float, int, int]] = []
|
candidates: list[tuple[float, int, int]] = []
|
||||||
for c_idx, content in enumerate(content_clusters):
|
for c_idx, content in enumerate(content_clusters):
|
||||||
want_below = content.boxclass == "picture" # figure 标题在下
|
cluster_kind = "figure" if content.boxclass == "picture" else "table"
|
||||||
want_kind = "figure" if want_below else "table"
|
|
||||||
for b_idx, cap in enumerate(caption_blocks):
|
for b_idx, cap in enumerate(caption_blocks):
|
||||||
if cap.kind != want_kind:
|
if cap.kind != cluster_kind: # 类型过滤:防 figure↔table 串台
|
||||||
continue
|
continue
|
||||||
cx0, cy0, cx1, cy1 = cap.bbox
|
cx0, cy0, cx1, cy1 = cap.bbox
|
||||||
h_overlap = min(content.x1, cx1) - max(content.x0, cx0)
|
h_overlap = min(content.x1, cx1) - max(content.x0, cx0)
|
||||||
@@ -185,24 +198,87 @@ def _pair_caption_blocks(
|
|||||||
if min_width <= 0 or h_overlap < min_width * 0.25:
|
if min_width <= 0 or h_overlap < min_width * 0.25:
|
||||||
continue
|
continue
|
||||||
if cy1 <= content.y0: # 标题在内容上方
|
if cy1 <= content.y0: # 标题在内容上方
|
||||||
side_below, v_gap = False, content.y0 - cy1
|
v_gap = content.y0 - cy1
|
||||||
elif cy0 >= content.y1: # 标题在内容下方
|
elif cy0 >= content.y1: # 标题在内容下方
|
||||||
side_below, v_gap = True, cy0 - content.y1
|
v_gap = cy0 - content.y1
|
||||||
else:
|
else:
|
||||||
continue # 重叠,跳过
|
continue # 重叠,跳过
|
||||||
if v_gap > _CAPTION_MATCH_DISTANCE:
|
if v_gap > _CAPTION_MATCH_DISTANCE:
|
||||||
continue
|
continue
|
||||||
penalty = 0.0 if side_below == want_below else _CAPTION_WRONG_SIDE_PENALTY
|
candidates.append((v_gap, c_idx, b_idx))
|
||||||
candidates.append((v_gap + penalty, c_idx, b_idx))
|
|
||||||
|
|
||||||
matches: dict[int, _CaptionBlock] = {}
|
# cluster 唯一归属最近的 caption;caption 可被多个 cluster 共享(复合图/子图)
|
||||||
used: set[int] = set()
|
cluster_to_caption: dict[int, int] = {}
|
||||||
for _score, c_idx, b_idx in sorted(candidates):
|
for _score, c_idx, b_idx in sorted(candidates):
|
||||||
if c_idx in matches or b_idx in used:
|
if c_idx in cluster_to_caption:
|
||||||
continue
|
continue
|
||||||
matches[c_idx] = caption_blocks[b_idx]
|
cluster_to_caption[c_idx] = b_idx
|
||||||
used.add(b_idx)
|
|
||||||
return matches
|
# 聚合 caption → clusters,按页面位置排序保证稳定的合并/渲染顺序
|
||||||
|
caption_to_clusters: dict[int, list[int]] = {}
|
||||||
|
for c_idx, b_idx in cluster_to_caption.items():
|
||||||
|
caption_to_clusters.setdefault(b_idx, []).append(c_idx)
|
||||||
|
for b_idx in caption_to_clusters:
|
||||||
|
caption_to_clusters[b_idx].sort(
|
||||||
|
key=lambda i: (content_clusters[i].y0, content_clusters[i].x0)
|
||||||
|
)
|
||||||
|
return caption_to_clusters
|
||||||
|
|
||||||
|
|
||||||
|
def _absorb_stragglers(
|
||||||
|
clusters: list[_BoxCluster],
|
||||||
|
caption_blocks: list[_CaptionBlock],
|
||||||
|
caption_matches: dict[int, list[int]],
|
||||||
|
) -> None:
|
||||||
|
"""把配不到 caption、却紧邻某已配 cluster 的同类型游离碎片并入该 cluster 所属 caption。
|
||||||
|
|
||||||
|
多面板图常把主 caption 只放在最下方,离 caption 过远(> _CAPTION_MATCH_DISTANCE)
|
||||||
|
的上方面板会被正常配对漏掉(如 Figure 4 的 (a) 子图距主标题 270pt)。这里把它们
|
||||||
|
并入紧邻的、已归属某 caption 的同类型 cluster,两道护栏避免误并:
|
||||||
|
|
||||||
|
1. 垂直紧邻(≤ _ABSORB_GAP)且水平重叠(同一列)——排除不同列的无关图;
|
||||||
|
2. 两者之间不得夹其他 caption ——有 caption 即另一张图/表的边界(如 Table 7/8
|
||||||
|
之间夹 Table 7 标题),不并,从而不破坏表格的独立配对。
|
||||||
|
"""
|
||||||
|
# cluster_idx → caption_idx(已配 cluster 的反向索引)
|
||||||
|
cluster_to_cap: dict[int, int] = {}
|
||||||
|
for cap_idx, idxs in caption_matches.items():
|
||||||
|
for c_idx in idxs:
|
||||||
|
cluster_to_cap[c_idx] = cap_idx
|
||||||
|
|
||||||
|
# 多趟扫描直到稳定:游离面板可能链式排列(上图→中图→下图→caption),
|
||||||
|
# 单趟只能吸收紧邻已配 cluster 的那一层;中图被并入后才轮到上图,故需重复。
|
||||||
|
changed = True
|
||||||
|
while changed:
|
||||||
|
changed = False
|
||||||
|
for u_idx, u in enumerate(clusters):
|
||||||
|
if u_idx in cluster_to_cap:
|
||||||
|
continue # 已配,无需吸收
|
||||||
|
best_cap: int | None = None
|
||||||
|
best_gap: float | None = None
|
||||||
|
for p_idx, cap_idx in cluster_to_cap.items():
|
||||||
|
p = clusters[p_idx]
|
||||||
|
if p.boxclass != u.boxclass:
|
||||||
|
continue # 类型不同(figure vs table)不并
|
||||||
|
gap = max(0.0, max(u.y0, p.y0) - min(u.y1, p.y1))
|
||||||
|
if gap > _ABSORB_GAP:
|
||||||
|
continue # 护栏 1:垂直不紧邻
|
||||||
|
if min(u.x1, p.x1) - max(u.x0, p.x0) <= 0:
|
||||||
|
continue # 护栏 1:水平不重叠,非同一列
|
||||||
|
# 护栏 2:u 与 p 之间的垂直区间不得夹任何 caption
|
||||||
|
between_lo = min(u.y1, p.y1)
|
||||||
|
between_hi = max(u.y0, p.y0)
|
||||||
|
if any(
|
||||||
|
not (cb.bbox[3] < between_lo or cb.bbox[1] > between_hi)
|
||||||
|
for cb in caption_blocks
|
||||||
|
):
|
||||||
|
continue
|
||||||
|
if best_gap is None or gap < best_gap:
|
||||||
|
best_gap, best_cap = gap, cap_idx
|
||||||
|
if best_cap is not None:
|
||||||
|
caption_matches[best_cap].append(u_idx)
|
||||||
|
cluster_to_cap[u_idx] = best_cap # 标记已并入,供后续趟链式吸收
|
||||||
|
changed = True
|
||||||
|
|
||||||
|
|
||||||
# ── Phase 1: 检测 + 渲染 ──────────────────────────────────────────────
|
# ── Phase 1: 检测 + 渲染 ──────────────────────────────────────────────
|
||||||
@@ -259,6 +335,7 @@ def _process_page(
|
|||||||
"""处理单页:检测内容 box → 文本定位 caption → 只渲染配到标题的。
|
"""处理单页:检测内容 box → 文本定位 caption → 只渲染配到标题的。
|
||||||
|
|
||||||
配到 Figure/Table caption 的 box 用 caption 自带 ID 命名(figure_3.jpg);
|
配到 Figure/Table caption 的 box 用 caption 自带 ID 命名(figure_3.jpg);
|
||||||
|
同一 caption 的多个 cluster(复合图/子图被 DocLayout 切散)合并 bbox 整张截取;
|
||||||
没配到标题的(Algorithm 伪代码、无编号附录表、误检碎片)一律过滤,不输出。
|
没配到标题的(Algorithm 伪代码、无编号附录表、误检碎片)一律过滤,不输出。
|
||||||
"""
|
"""
|
||||||
page = doc[page_idx]
|
page = doc[page_idx]
|
||||||
@@ -280,38 +357,46 @@ def _process_page(
|
|||||||
# 聚类:将同一 figure/table 的碎片 box 合并;用 PDF 文本定位 caption
|
# 聚类:将同一 figure/table 的碎片 box 合并;用 PDF 文本定位 caption
|
||||||
clusters = _cluster_boxes(raw_boxes)
|
clusters = _cluster_boxes(raw_boxes)
|
||||||
caption_blocks = _find_caption_blocks(page)
|
caption_blocks = _find_caption_blocks(page)
|
||||||
|
# caption_idx → [cluster_idx, ...];一个 caption 可含多个子图 cluster(复合图)
|
||||||
caption_matches = _pair_caption_blocks(clusters, caption_blocks)
|
caption_matches = _pair_caption_blocks(clusters, caption_blocks)
|
||||||
|
# 把距主 caption 过远的游离面板(如多面板图上方的 (a) 子图)并入紧邻的已配 cluster
|
||||||
|
_absorb_stragglers(clusters, caption_blocks, caption_matches)
|
||||||
|
|
||||||
extracted = 0
|
extracted = 0
|
||||||
for cluster_idx, cluster in enumerate(clusters):
|
for cap_idx, cluster_indices in caption_matches.items():
|
||||||
cap_match = caption_matches.get(cluster_idx)
|
cap = caption_blocks[cap_idx]
|
||||||
if cap_match is None:
|
if cap.id in seen_labels:
|
||||||
continue # 无 Figure/Table 标题 → 过滤(Algorithm、无编号表、误检碎片)
|
|
||||||
if cap_match.id in seen_labels:
|
|
||||||
continue # 同一图表被 DocLayout 切成多块重复检测,跳过后续
|
continue # 同一图表被 DocLayout 切成多块重复检测,跳过后续
|
||||||
seen_labels.add(cap_match.id)
|
seen_labels.add(cap.id)
|
||||||
|
|
||||||
filename = f"{cap_match.id.replace(' ', '_').lower()}.jpg"
|
# 同一 caption 的所有 cluster 合并 bbox,复合图/子图整张截取
|
||||||
|
members = [clusters[i] for i in cluster_indices]
|
||||||
|
merged = _BoxCluster(members)
|
||||||
|
|
||||||
|
filename = f"{cap.id.replace(' ', '_').lower()}.jpg"
|
||||||
if not _render_box(
|
if not _render_box(
|
||||||
page,
|
page,
|
||||||
cluster,
|
merged,
|
||||||
images_dest,
|
images_dest,
|
||||||
filename,
|
filename,
|
||||||
cap_match.kind,
|
cap.kind,
|
||||||
page_num,
|
page_num,
|
||||||
caption_bbox=cap_match.bbox,
|
caption_bbox=cap.bbox,
|
||||||
):
|
):
|
||||||
continue
|
continue
|
||||||
|
|
||||||
manifest[filename] = {
|
info = {
|
||||||
"page": page_num,
|
"page": page_num,
|
||||||
"type": cap_match.kind,
|
"type": cap.kind,
|
||||||
"label": cap_match.id,
|
"label": cap.id,
|
||||||
"box": _cluster_to_box(cluster),
|
"box": _cluster_to_box(merged),
|
||||||
"caption_text": cap_match.text[:500],
|
"caption_text": cap.text[:500],
|
||||||
"caption_box": cap_match.bbox,
|
"caption_box": cap.bbox,
|
||||||
"caption_source": "text",
|
"caption_source": "text",
|
||||||
}
|
}
|
||||||
|
if len(members) > 1:
|
||||||
|
info["subfigure_count"] = len(members)
|
||||||
|
manifest[filename] = info
|
||||||
extracted += 1
|
extracted += 1
|
||||||
|
|
||||||
return extracted
|
return extracted
|
||||||
@@ -458,19 +543,31 @@ def link_figures_with_images(
|
|||||||
if not unmatched:
|
if not unmatched:
|
||||||
return figures
|
return figures
|
||||||
|
|
||||||
|
# 已被策略 1(精确匹配)占用的图片不参与兜底,否则会把已正确归属的图复用给
|
||||||
|
# 别的条目(如缺失的 Table 4 误链到 Table 1 的截图)。
|
||||||
|
assigned_urls = {f["image_url"] for f in figures if f.get("image_url")}
|
||||||
|
|
||||||
# 按类型分流:Figure vs Table
|
# 按类型分流:Figure vs Table
|
||||||
fig_type_unmatched = [f for f in unmatched if _is_figure_type(f.get("id", ""))]
|
fig_type_unmatched = [f for f in unmatched if _is_figure_type(f.get("id", ""))]
|
||||||
table_type_unmatched = [
|
table_type_unmatched = [
|
||||||
f for f in unmatched if not _is_figure_type(f.get("id", ""))
|
f for f in unmatched if not _is_figure_type(f.get("id", ""))
|
||||||
]
|
]
|
||||||
|
|
||||||
# 提取的图片按类型分流,按文件名中的编号排序
|
# 剩余未占用图片按类型分流,按文件名中的编号排序
|
||||||
fig_images = sorted(
|
fig_images = sorted(
|
||||||
[img for img in images if "table" not in img["name"].lower()],
|
[
|
||||||
|
img
|
||||||
|
for img in images
|
||||||
|
if "table" not in img["name"].lower() and img["url"] not in assigned_urls
|
||||||
|
],
|
||||||
key=lambda img: _image_sort_key(img["name"]),
|
key=lambda img: _image_sort_key(img["name"]),
|
||||||
)
|
)
|
||||||
table_images = sorted(
|
table_images = sorted(
|
||||||
[img for img in images if "table" in img["name"].lower()],
|
[
|
||||||
|
img
|
||||||
|
for img in images
|
||||||
|
if "table" in img["name"].lower() and img["url"] not in assigned_urls
|
||||||
|
],
|
||||||
key=lambda img: _image_sort_key(img["name"]),
|
key=lambda img: _image_sort_key(img["name"]),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
+3
-2
@@ -93,13 +93,14 @@ def release_lock(db, lock) -> None:
|
|||||||
|
|
||||||
|
|
||||||
def make_http_client(
|
def make_http_client(
|
||||||
*, sync: bool = False, follow_redirects: bool = False, **kwargs
|
*, sync: bool = False, follow_redirects: bool = False, use_proxy: bool = True, **kwargs
|
||||||
) -> httpx.AsyncClient | httpx.Client:
|
) -> httpx.AsyncClient | httpx.Client:
|
||||||
"""创建带 proxy 和默认配置的 httpx 客户端。
|
"""创建带 proxy 和默认配置的 httpx 客户端。
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
sync: True 返回同步 Client,False 返回 AsyncClient
|
sync: True 返回同步 Client,False 返回 AsyncClient
|
||||||
follow_redirects: 是否跟随重定向
|
follow_redirects: 是否跟随重定向
|
||||||
|
use_proxy: 是否使用 HF_PROXY 代理(默认 True;国内 API 如 embedder 应传 False)
|
||||||
**kwargs: 覆盖默认参数
|
**kwargs: 覆盖默认参数
|
||||||
"""
|
"""
|
||||||
defaults: dict = {
|
defaults: dict = {
|
||||||
@@ -107,7 +108,7 @@ def make_http_client(
|
|||||||
"headers": {"User-Agent": settings.HTTP_USER_AGENT},
|
"headers": {"User-Agent": settings.HTTP_USER_AGENT},
|
||||||
"follow_redirects": follow_redirects,
|
"follow_redirects": follow_redirects,
|
||||||
}
|
}
|
||||||
if settings.http_proxy:
|
if use_proxy and settings.http_proxy:
|
||||||
defaults["transport"] = (
|
defaults["transport"] = (
|
||||||
httpx.HTTPTransport(proxy=settings.http_proxy)
|
httpx.HTTPTransport(proxy=settings.http_proxy)
|
||||||
if sync
|
if sync
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
from unittest.mock import MagicMock
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
import pymupdf
|
import pymupdf
|
||||||
@@ -168,3 +169,286 @@ def test_process_page_filters_uncaptioned(tmp_path):
|
|||||||
|
|
||||||
assert extracted == 0
|
assert extracted == 0
|
||||||
assert manifest == {}
|
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"]
|
||||||
|
|||||||
Reference in New Issue
Block a user