mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 06:50:53 +00:00
revert: restore original Phase 2 figure/table matching logic
This commit is contained in:
Research Assistant
parent
ac3c0c4b8e
commit
0a2988b89f
4 changed files with 78 additions and 543 deletions
|
|
@ -51,31 +51,11 @@ def _centroid_y(bbox: list[float]) -> float:
|
|||
return (bbox[1] + bbox[3]) / 2
|
||||
|
||||
|
||||
def _compute_text_confidence(text: str, bbox: list[float], profile_legends: list[dict]) -> bool:
|
||||
"""Check if a block's text/geometry is similar to known formal legends."""
|
||||
if not text or not profile_legends:
|
||||
return False
|
||||
if len(text) > 200:
|
||||
return False
|
||||
if any(w in text.lower() for w in [" is ", " are ", " was ", " were "]):
|
||||
return False
|
||||
if profile_legends:
|
||||
widths = [
|
||||
l.get("bbox", [0, 0, 0, 0])[2] - l.get("bbox", [0, 0, 0, 0])[0]
|
||||
for l in profile_legends
|
||||
if l.get("bbox")
|
||||
]
|
||||
if widths:
|
||||
avg_width = sum(widths) / len(widths)
|
||||
bbox_width = bbox[2] - bbox[0]
|
||||
if abs(bbox_width - avg_width) / max(1, avg_width) < 0.5:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def build_figure_inventory(structured_blocks: list[dict]) -> dict[str, Any]:
|
||||
legends: list[dict] = []
|
||||
assets: list[dict] = []
|
||||
unmatched_assets: list[dict] = []
|
||||
matched_figures: list[dict] = []
|
||||
|
||||
for block in structured_blocks:
|
||||
role = block.get("role", "")
|
||||
|
|
@ -88,109 +68,84 @@ def build_figure_inventory(structured_blocks: list[dict]) -> dict[str, Any]:
|
|||
if raw_label in {"image", "chart", "figure_title", "figure"} or not raw_label:
|
||||
assets.append(block)
|
||||
|
||||
# --- Step 1: Classify each legend with confidence level ---
|
||||
# All legends are included; the classification affects confidence scoring
|
||||
legend_classification: dict[int, str] = {} # legend index -> "formal" | "candidate" | "fallback"
|
||||
for i, leg in enumerate(legends):
|
||||
text = leg.get("text", "")
|
||||
bbox = leg.get("bbox", [0, 0, 0, 0])
|
||||
if _extract_figure_number(text) is not None:
|
||||
legend_classification[i] = "formal"
|
||||
elif _compute_text_confidence(text, bbox, [legends[j] for j in range(i) if legend_classification.get(j) == "formal"]):
|
||||
legend_classification[i] = "candidate"
|
||||
else:
|
||||
legend_classification[i] = "fallback"
|
||||
|
||||
# --- Step 2: Multi-signal matching ---
|
||||
used_asset_indices: set[int] = set()
|
||||
matched_figures: list[dict] = []
|
||||
unmatched_legends: list[dict] = []
|
||||
|
||||
for i_leg, legend in enumerate(legends):
|
||||
for legend in legends:
|
||||
legend_page = legend.get("page", 0)
|
||||
legend_bbox = legend.get("bbox", [0, 0, 0, 0])
|
||||
legend_text = legend.get("text", "")
|
||||
fig_num = _extract_figure_number(legend_text)
|
||||
cls = legend_classification.get(i_leg, "fallback")
|
||||
is_formal = cls == "formal"
|
||||
|
||||
candidate_pages = [legend_page, legend_page + 1, legend_page - 1]
|
||||
matched_assets = []
|
||||
|
||||
matched_assets = []
|
||||
for page in candidate_pages:
|
||||
if page < 1:
|
||||
continue
|
||||
page_assets = [
|
||||
(i, a) for i, a in enumerate(assets)
|
||||
if a.get("page", 0) == page and i not in used_asset_indices
|
||||
(i, a) for i, a in enumerate(assets) if a.get("page", 0) == page and i not in used_asset_indices
|
||||
]
|
||||
if not page_assets:
|
||||
continue
|
||||
|
||||
scored = []
|
||||
candidates_for_page: list[dict] = []
|
||||
for i, asset in page_assets:
|
||||
asset_bbox = asset.get("bbox", [0, 0, 0, 0])
|
||||
overlap = _compute_overlap_score(legend_bbox, asset_bbox)
|
||||
dist_y = abs(_centroid_y(legend_bbox) - _centroid_y(asset_bbox))
|
||||
dir_bonus = 1.0 if _centroid_y(asset_bbox) < _centroid_y(legend_bbox) else 0.5
|
||||
same_page = 2.0 if page == legend_page else 0.0
|
||||
# Horizontal alignment
|
||||
legend_cx = (legend_bbox[0] + legend_bbox[2]) / 2
|
||||
asset_cx = (asset_bbox[0] + asset_bbox[2]) / 2
|
||||
h_align = max(0, 1.0 - abs(legend_cx - asset_cx) / max(1, legend_bbox[2] - legend_bbox[0]))
|
||||
# Asset size bonus
|
||||
a_w = asset_bbox[2] - asset_bbox[0]
|
||||
a_h = asset_bbox[3] - asset_bbox[1]
|
||||
size_score = min(2.0, (a_w * a_h) / 50000)
|
||||
score = overlap * 10 + dir_bonus + same_page + h_align * 2 + size_score * 0.5 - dist_y * 0.01
|
||||
scored.append({"asset_index": i, "score": score, "overlap": overlap})
|
||||
direction_bonus = 1.0 if _centroid_y(asset_bbox) < _centroid_y(legend_bbox) else 0.5
|
||||
same_page_bonus = 2.0 if page == legend_page else 0.0
|
||||
score = overlap * 10 + direction_bonus + same_page_bonus - dist_y * 0.01
|
||||
candidates_for_page.append(
|
||||
{
|
||||
"asset_index": i,
|
||||
"asset": asset,
|
||||
"score": score,
|
||||
"overlap": overlap,
|
||||
"distance_y": dist_y,
|
||||
"same_page": page == legend_page,
|
||||
}
|
||||
)
|
||||
|
||||
scored.sort(key=lambda c: c["score"], reverse=True)
|
||||
for c in scored:
|
||||
if c["score"] <= -5:
|
||||
continue
|
||||
candidates_for_page.sort(key=lambda c: c["score"], reverse=True)
|
||||
|
||||
for candidate in candidates_for_page:
|
||||
if len(matched_assets) >= 3:
|
||||
break
|
||||
matched_assets.append(assets[c["asset_index"]])
|
||||
used_asset_indices.add(c["asset_index"])
|
||||
if candidate["score"] > -5:
|
||||
matched_assets.append(candidate["asset"])
|
||||
used_asset_indices.add(candidate["asset_index"])
|
||||
|
||||
if matched_assets:
|
||||
break
|
||||
|
||||
# --- Step 3: Classify outcome ---
|
||||
flags: list[str] = []
|
||||
confidence = 0.85 if is_formal else 0.5
|
||||
if not matched_assets:
|
||||
flags.append("legend_only")
|
||||
confidence = 0.4 if is_formal else 0.2
|
||||
unmatched_legends.append(legend)
|
||||
elif confidence < 0.6:
|
||||
flags.append("low_confidence_match")
|
||||
confidence = 0.4
|
||||
matched_figures.append(
|
||||
{
|
||||
"legend_block_id": legend.get("block_id", ""),
|
||||
"page": legend_page,
|
||||
"text": legend_text,
|
||||
"figure_number": fig_num,
|
||||
"matched_assets": [
|
||||
{
|
||||
"block_id": a.get("block_id", ""),
|
||||
"bbox": a.get("bbox", [0, 0, 0, 0]),
|
||||
}
|
||||
for a in matched_assets
|
||||
],
|
||||
"confidence": 0.85 if matched_assets else 0.4,
|
||||
"flags": [] if matched_assets else ["legend_only"],
|
||||
}
|
||||
)
|
||||
|
||||
matched_figures.append({
|
||||
"legend_block_id": legend.get("block_id", ""),
|
||||
"page": legend_page,
|
||||
"text": legend_text,
|
||||
"figure_number": fig_num,
|
||||
"matched_assets": [
|
||||
{"block_id": a.get("block_id", ""), "bbox": a.get("bbox", [0, 0, 0, 0])}
|
||||
for a in matched_assets
|
||||
],
|
||||
"confidence": confidence,
|
||||
"flags": flags,
|
||||
})
|
||||
|
||||
unmatched_assets_list = [
|
||||
asset for i, asset in enumerate(assets) if i not in used_asset_indices
|
||||
]
|
||||
for i, asset in enumerate(assets):
|
||||
if i not in used_asset_indices:
|
||||
unmatched_assets.append(asset)
|
||||
|
||||
return {
|
||||
"figure_legends": legends,
|
||||
"figure_assets": assets,
|
||||
"matched_figures": matched_figures,
|
||||
"unmatched_legends": unmatched_legends,
|
||||
"unmatched_assets": unmatched_assets_list,
|
||||
"unmatched_legends": [],
|
||||
"unmatched_assets": unmatched_assets,
|
||||
"official_figure_count": len(matched_figures),
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -6,13 +6,12 @@ from typing import Any
|
|||
|
||||
from paperforge.core.io import write_json
|
||||
|
||||
|
||||
_TABLE_PREFIX_PATTERN = re.compile(
|
||||
r"^(?:Table|Supplementary\s+Table|Extended\s+Data\s+Table)\s+(\d+(?:\.\d+)?)",
|
||||
flags=re.IGNORECASE,
|
||||
)
|
||||
|
||||
_CONTINUATION_PATTERN = re.compile(r"\(Continued\)|\(cont\.\)|\(cont\)", re.IGNORECASE)
|
||||
|
||||
|
||||
def _extract_table_number(text: str) -> int | None:
|
||||
m = _TABLE_PREFIX_PATTERN.search(text)
|
||||
|
|
@ -24,31 +23,12 @@ def _extract_table_number(text: str) -> int | None:
|
|||
return None
|
||||
|
||||
|
||||
def _compute_overlap_score(a_bbox: list[float], b_bbox: list[float]) -> float:
|
||||
if not a_bbox or not b_bbox or len(a_bbox) < 4 or len(b_bbox) < 4:
|
||||
return 0.0
|
||||
ax1, ay1, ax2, ay2 = a_bbox
|
||||
bx1, by1, bx2, by2 = b_bbox
|
||||
ix1 = max(ax1, bx1)
|
||||
iy1 = max(ay1, by1)
|
||||
ix2 = min(ax2, bx2)
|
||||
iy2 = min(ay2, by2)
|
||||
iw = max(0, ix2 - ix1)
|
||||
ih = max(0, iy2 - iy1)
|
||||
intersection = iw * ih
|
||||
if intersection == 0:
|
||||
return 0.0
|
||||
a_area = (ax2 - ax1) * (ay2 - ay1)
|
||||
b_area = (bx2 - bx1) * (by2 - by1)
|
||||
union = a_area + b_area - intersection
|
||||
if union == 0:
|
||||
return 0.0
|
||||
return intersection / union
|
||||
|
||||
|
||||
def build_table_inventory(structured_blocks: list[dict]) -> dict[str, Any]:
|
||||
tables: list[dict] = []
|
||||
captions: list[dict] = []
|
||||
assets: list[dict] = []
|
||||
unmatched_captions: list[dict] = []
|
||||
unmatched_assets: list[dict] = []
|
||||
|
||||
for block in structured_blocks:
|
||||
role = block.get("role", "")
|
||||
|
|
@ -61,72 +41,30 @@ def build_table_inventory(structured_blocks: list[dict]) -> dict[str, Any]:
|
|||
assets.append(block)
|
||||
|
||||
used_asset_indices: set[int] = set()
|
||||
tables: list[dict] = []
|
||||
unmatched_captions: list[dict] = []
|
||||
|
||||
for caption in captions:
|
||||
caption_page = caption.get("page", 0)
|
||||
caption_text = caption.get("text", "") or ""
|
||||
caption_bbox = caption.get("bbox", [0, 0, 0, 0])
|
||||
caption_text = caption.get("text", "")
|
||||
table_num = _extract_table_number(caption_text)
|
||||
is_continuation = bool(_CONTINUATION_PATTERN.search(caption_text))
|
||||
|
||||
# Continuations should only match same-page assets
|
||||
candidate_pages = [caption_page]
|
||||
if not is_continuation:
|
||||
candidate_pages.extend([caption_page + 1, caption_page - 1])
|
||||
candidate_pages = [caption_page, caption_page + 1, caption_page - 1]
|
||||
matched_asset = None
|
||||
|
||||
best_match: tuple[int | None, float] = (None, -999)
|
||||
for page in candidate_pages:
|
||||
if page < 1:
|
||||
continue
|
||||
page_assets = [
|
||||
(i, a) for i, a in enumerate(assets)
|
||||
if a.get("page", 0) == page and i not in used_asset_indices
|
||||
]
|
||||
for i, asset in page_assets:
|
||||
asset_bbox = asset.get("bbox", [0, 0, 0, 0])
|
||||
overlap = _compute_overlap_score(caption_bbox, asset_bbox)
|
||||
if is_continuation:
|
||||
same_page = 3.0
|
||||
h_align = 0.0
|
||||
if len(caption_bbox) >= 4 and len(asset_bbox) >= 4:
|
||||
ccx = (caption_bbox[0] + caption_bbox[2]) / 2
|
||||
acx = (asset_bbox[0] + asset_bbox[2]) / 2
|
||||
h_align = max(0, 1.0 - abs(ccx - acx) / max(1, caption_bbox[2] - caption_bbox[0]))
|
||||
score = same_page + h_align * 3 + overlap * 5
|
||||
else:
|
||||
dist_y = 0.0
|
||||
if len(caption_bbox) >= 4 and len(asset_bbox) >= 4:
|
||||
dist_y = abs((caption_bbox[1] + caption_bbox[3]) / 2 - (asset_bbox[1] + asset_bbox[3]) / 2)
|
||||
same_page = 2.0 if page == caption_page else 0.0
|
||||
direction = 1.0 if (len(asset_bbox) >= 4 and len(caption_bbox) >= 4
|
||||
and (asset_bbox[1] + asset_bbox[3]) / 2 < (caption_bbox[1] + caption_bbox[3]) / 2) else 0.5
|
||||
h_align = 0.0
|
||||
if len(caption_bbox) >= 4 and len(asset_bbox) >= 4:
|
||||
ccx = (caption_bbox[0] + caption_bbox[2]) / 2
|
||||
acx = (asset_bbox[0] + asset_bbox[2]) / 2
|
||||
h_align = max(0, 1.0 - abs(ccx - acx) / max(1, caption_bbox[2] - caption_bbox[0]))
|
||||
a_w = asset_bbox[2] - asset_bbox[0] if len(asset_bbox) >= 4 else 0
|
||||
a_h = asset_bbox[3] - asset_bbox[1] if len(asset_bbox) >= 4 else 0
|
||||
size_score = min(2.0, (a_w * a_h) / 50000)
|
||||
score = overlap * 10 + same_page + direction + h_align * 2 + size_score * 0.5 - dist_y * 0.01
|
||||
if score > best_match[1]:
|
||||
best_match = (i, score)
|
||||
|
||||
matched_asset = assets[best_match[0]] if best_match[0] is not None else None
|
||||
if best_match[0] is not None:
|
||||
used_asset_indices.add(best_match[0])
|
||||
|
||||
if matched_asset is None:
|
||||
unmatched_captions.append(caption)
|
||||
for i, asset in enumerate(assets):
|
||||
if asset.get("page", 0) == page and i not in used_asset_indices:
|
||||
matched_asset = asset
|
||||
used_asset_indices.add(i)
|
||||
break
|
||||
if matched_asset:
|
||||
break
|
||||
|
||||
tables.append({
|
||||
"caption_block_id": caption.get("block_id", ""),
|
||||
"page": caption_page,
|
||||
"caption_text": caption_text,
|
||||
"table_number": table_num,
|
||||
"is_continuation": is_continuation,
|
||||
"asset_block_id": matched_asset.get("block_id", "") if matched_asset else "",
|
||||
"asset_bbox": matched_asset.get("bbox", [0, 0, 0, 0]) if matched_asset else [],
|
||||
"assistive_text": (matched_asset.get("text", "") or "") if matched_asset else "",
|
||||
|
|
@ -134,17 +72,23 @@ def build_table_inventory(structured_blocks: list[dict]) -> dict[str, Any]:
|
|||
"has_asset": matched_asset is not None,
|
||||
})
|
||||
|
||||
unmatched_assets_list = [
|
||||
a for i, a in enumerate(assets) if i not in used_asset_indices
|
||||
]
|
||||
for caption in captions:
|
||||
if not any(
|
||||
t["caption_block_id"] == caption.get("block_id", "")
|
||||
for t in tables
|
||||
if t["has_asset"]
|
||||
):
|
||||
unmatched_captions.append(caption)
|
||||
|
||||
official_count = len([t for t in tables if t["has_asset"] and not t["is_continuation"]])
|
||||
for i, asset in enumerate(assets):
|
||||
if i not in used_asset_indices:
|
||||
unmatched_assets.append(asset)
|
||||
|
||||
return {
|
||||
"tables": tables,
|
||||
"unmatched_captions": unmatched_captions,
|
||||
"unmatched_assets": unmatched_assets_list,
|
||||
"official_table_count": official_count,
|
||||
"unmatched_assets": unmatched_assets,
|
||||
"official_table_count": len([t for t in tables if t["has_asset"]]),
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -74,149 +74,3 @@ def test_unmatched_assets_are_preserved() -> None:
|
|||
|
||||
assert inventory["official_figure_count"] == 0
|
||||
assert len(inventory["unmatched_assets"]) == 1
|
||||
|
||||
|
||||
def test_extract_figure_number_basic() -> None:
|
||||
from paperforge.worker.ocr_figures import _extract_figure_number
|
||||
assert _extract_figure_number("Figure 1. Caption") == 1
|
||||
|
||||
|
||||
def test_extract_figure_number_fig_dot() -> None:
|
||||
from paperforge.worker.ocr_figures import _extract_figure_number
|
||||
assert _extract_figure_number("Fig. 2. Test") == 2
|
||||
|
||||
|
||||
def test_extract_figure_number_supplementary() -> None:
|
||||
from paperforge.worker.ocr_figures import _extract_figure_number
|
||||
assert _extract_figure_number("Supplementary Fig. S3") == 3
|
||||
|
||||
|
||||
def test_extract_figure_number_extended_data() -> None:
|
||||
from paperforge.worker.ocr_figures import _extract_figure_number
|
||||
assert _extract_figure_number("Extended Data Fig. 4.") == 4
|
||||
|
||||
|
||||
def test_extract_figure_number_decimal_truncated() -> None:
|
||||
from paperforge.worker.ocr_figures import _extract_figure_number
|
||||
result = _extract_figure_number("Figure 1.2. Magnified view")
|
||||
assert result == 1 or result == 1.2
|
||||
|
||||
|
||||
def test_extract_figure_number_none() -> None:
|
||||
from paperforge.worker.ocr_figures import _extract_figure_number
|
||||
assert _extract_figure_number("Some random text") is None
|
||||
|
||||
|
||||
def test_extract_figure_number_multiline() -> None:
|
||||
from paperforge.worker.ocr_figures import _extract_figure_number
|
||||
assert _extract_figure_number("Figure 3.\nDescription continues") == 3
|
||||
|
||||
|
||||
def test_formal_legend_detection_explicit_figure_prefix() -> None:
|
||||
from paperforge.worker.ocr_figures import build_figure_inventory
|
||||
|
||||
structured_blocks = [
|
||||
{
|
||||
"paper_id": "K001",
|
||||
"page": 1,
|
||||
"block_id": "p1_b1",
|
||||
"role": "figure_caption",
|
||||
"text": "Figure 1. This is a formal legend with plenty of descriptive text that explains the figure contents in detail.",
|
||||
"bbox": [50, 400, 550, 450],
|
||||
},
|
||||
{
|
||||
"paper_id": "K001",
|
||||
"page": 1,
|
||||
"block_id": "p1_b2",
|
||||
"role": "figure_asset",
|
||||
"text": "",
|
||||
"bbox": [50, 50, 550, 380],
|
||||
},
|
||||
]
|
||||
|
||||
inventory = build_figure_inventory(structured_blocks)
|
||||
|
||||
assert len(inventory["matched_figures"]) == 1
|
||||
assert inventory["matched_figures"][0]["figure_number"] == 1
|
||||
assert len(inventory["matched_figures"][0]["matched_assets"]) == 1
|
||||
assert inventory["matched_figures"][0]["confidence"] == 0.85
|
||||
|
||||
|
||||
def test_candidate_legend_geometry_match() -> None:
|
||||
from paperforge.worker.ocr_figures import build_figure_inventory
|
||||
|
||||
structured_blocks = [
|
||||
{
|
||||
"paper_id": "K001",
|
||||
"page": 1,
|
||||
"block_id": "p1_b1",
|
||||
"role": "figure_caption",
|
||||
"text": "Figure 1. Formal legend that establishes a width profile.",
|
||||
"bbox": [50, 420, 550, 460],
|
||||
},
|
||||
{
|
||||
"paper_id": "K001",
|
||||
"page": 1,
|
||||
"block_id": "p1_b2",
|
||||
"role": "figure_caption",
|
||||
"text": "No figure prefix but short and profile-matched",
|
||||
"bbox": [60, 350, 540, 380],
|
||||
},
|
||||
{
|
||||
"paper_id": "K001",
|
||||
"page": 1,
|
||||
"block_id": "p1_b3",
|
||||
"role": "figure_asset",
|
||||
"text": "",
|
||||
"bbox": [60, 50, 540, 330],
|
||||
},
|
||||
]
|
||||
|
||||
inventory = build_figure_inventory(structured_blocks)
|
||||
|
||||
assert len(inventory["matched_figures"]) == 2
|
||||
match_texts = [m["text"] for m in inventory["matched_figures"]]
|
||||
assert any("No figure prefix" in t for t in match_texts)
|
||||
|
||||
|
||||
def test_legend_only_figure_no_asset_match() -> None:
|
||||
from paperforge.worker.ocr_figures import build_figure_inventory
|
||||
|
||||
structured_blocks = [
|
||||
{
|
||||
"paper_id": "K001",
|
||||
"page": 2,
|
||||
"block_id": "p2_b1",
|
||||
"role": "figure_caption",
|
||||
"text": "Figure 2. This caption has no matching asset on the same page.",
|
||||
"bbox": [50, 700, 550, 750],
|
||||
},
|
||||
]
|
||||
|
||||
inventory = build_figure_inventory(structured_blocks)
|
||||
|
||||
assert len(inventory["matched_figures"]) == 1
|
||||
assert inventory["matched_figures"][0]["figure_number"] == 2
|
||||
assert len(inventory["matched_figures"][0]["matched_assets"]) == 0
|
||||
assert "legend_only" in inventory["matched_figures"][0]["flags"]
|
||||
assert inventory["matched_figures"][0]["confidence"] == 0.4
|
||||
|
||||
|
||||
def test_unmatched_legends_populated() -> None:
|
||||
from paperforge.worker.ocr_figures import build_figure_inventory
|
||||
|
||||
structured_blocks = [
|
||||
{
|
||||
"paper_id": "K001",
|
||||
"page": 3,
|
||||
"block_id": "p3_b1",
|
||||
"role": "figure_caption",
|
||||
"text": "Figure 3. Caption with no figure asset at all.",
|
||||
"bbox": [50, 700, 550, 750],
|
||||
},
|
||||
]
|
||||
|
||||
inventory = build_figure_inventory(structured_blocks)
|
||||
|
||||
assert len(inventory["unmatched_legends"]) == 1
|
||||
assert inventory["unmatched_legends"][0]["block_id"] == "p3_b1"
|
||||
|
|
|
|||
|
|
@ -1,4 +1,8 @@
|
|||
"""Task 5 tests for hardened table matching."""
|
||||
"""Phase 2 contract tests for table inventory.
|
||||
|
||||
paperforge.worker.ocr_tables does not exist yet -- tests will fail until
|
||||
Task 7 implements the module.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
|
|
@ -59,225 +63,3 @@ def test_table_without_asset_is_tracked_as_unmatched_caption() -> None:
|
|||
|
||||
assert inventory["official_table_count"] == 0
|
||||
assert len(inventory["unmatched_captions"]) == 1
|
||||
|
||||
|
||||
def test_continuation_table_matches_same_page_asset() -> None:
|
||||
from paperforge.worker.ocr_tables import build_table_inventory
|
||||
|
||||
structured_blocks = [
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 10,
|
||||
"block_id": "p10_a1",
|
||||
"role": "table_asset",
|
||||
"text": "continued data",
|
||||
"bbox": [100, 100, 600, 400],
|
||||
},
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 10,
|
||||
"block_id": "p10_c1",
|
||||
"role": "table_caption",
|
||||
"text": "Table 1 (Continued)",
|
||||
"bbox": [100, 420, 600, 460],
|
||||
},
|
||||
]
|
||||
|
||||
inventory = build_table_inventory(structured_blocks)
|
||||
|
||||
assert len(inventory["tables"]) == 1
|
||||
t = inventory["tables"][0]
|
||||
assert t["is_continuation"] is True
|
||||
assert t["has_asset"] is True
|
||||
assert t["asset_block_id"] == "p10_a1"
|
||||
|
||||
|
||||
def test_continuation_does_not_increment_official_count() -> None:
|
||||
from paperforge.worker.ocr_tables import build_table_inventory
|
||||
|
||||
structured_blocks = [
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 5,
|
||||
"block_id": "p5_a1",
|
||||
"role": "table_asset",
|
||||
"text": "table data",
|
||||
"bbox": [100, 100, 600, 400],
|
||||
},
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 5,
|
||||
"block_id": "p5_c1",
|
||||
"role": "table_caption",
|
||||
"text": "Table 1. Main data",
|
||||
"bbox": [100, 420, 600, 460],
|
||||
},
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 6,
|
||||
"block_id": "p6_a1",
|
||||
"role": "table_asset",
|
||||
"text": "continued data",
|
||||
"bbox": [100, 100, 600, 300],
|
||||
},
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 6,
|
||||
"block_id": "p6_c1",
|
||||
"role": "table_caption",
|
||||
"text": "Table 1 (Continued)",
|
||||
"bbox": [100, 320, 600, 360],
|
||||
},
|
||||
]
|
||||
|
||||
inventory = build_table_inventory(structured_blocks)
|
||||
|
||||
assert inventory["official_table_count"] == 1
|
||||
normal = [t for t in inventory["tables"] if not t["is_continuation"]]
|
||||
continued = [t for t in inventory["tables"] if t["is_continuation"]]
|
||||
assert len(normal) == 1
|
||||
assert len(continued) == 1
|
||||
assert normal[0]["has_asset"] is True
|
||||
assert continued[0]["has_asset"] is True
|
||||
|
||||
|
||||
def test_continuation_without_asset_has_has_asset_false() -> None:
|
||||
from paperforge.worker.ocr_tables import build_table_inventory
|
||||
|
||||
structured_blocks = [
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 10,
|
||||
"block_id": "p10_c1",
|
||||
"role": "table_caption",
|
||||
"text": "Table 1 (Continued)",
|
||||
"bbox": [100, 100, 600, 140],
|
||||
},
|
||||
]
|
||||
|
||||
inventory = build_table_inventory(structured_blocks)
|
||||
|
||||
assert len(inventory["tables"]) == 1
|
||||
assert inventory["tables"][0]["is_continuation"] is True
|
||||
assert inventory["tables"][0]["has_asset"] is False
|
||||
assert len(inventory["unmatched_captions"]) == 1
|
||||
|
||||
|
||||
def test_multi_signal_scoring_prefers_better_asset() -> None:
|
||||
from paperforge.worker.ocr_tables import build_table_inventory
|
||||
|
||||
structured_blocks = [
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 3,
|
||||
"block_id": "p3_a1",
|
||||
"role": "table_asset",
|
||||
"text": "far table",
|
||||
"bbox": [50, 600, 300, 800],
|
||||
},
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 3,
|
||||
"block_id": "p3_a2",
|
||||
"role": "table_asset",
|
||||
"text": "near table",
|
||||
"bbox": [50, 50, 550, 300],
|
||||
},
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 3,
|
||||
"block_id": "p3_c1",
|
||||
"role": "table_caption",
|
||||
"text": "Table 3. Nearby data",
|
||||
"bbox": [50, 310, 550, 350],
|
||||
},
|
||||
]
|
||||
|
||||
inventory = build_table_inventory(structured_blocks)
|
||||
|
||||
assert inventory["official_table_count"] == 1
|
||||
t = inventory["tables"][0]
|
||||
assert t["asset_block_id"] == "p3_a2"
|
||||
|
||||
|
||||
def test_continuation_matches_only_same_page_not_adjacent() -> None:
|
||||
from paperforge.worker.ocr_tables import build_table_inventory
|
||||
|
||||
structured_blocks = [
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 4,
|
||||
"block_id": "p4_a1",
|
||||
"role": "table_asset",
|
||||
"text": "wrong page",
|
||||
"bbox": [100, 100, 600, 400],
|
||||
},
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 5,
|
||||
"block_id": "p5_a1",
|
||||
"role": "table_asset",
|
||||
"text": "same page",
|
||||
"bbox": [100, 100, 600, 400],
|
||||
},
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 5,
|
||||
"block_id": "p5_c1",
|
||||
"role": "table_caption",
|
||||
"text": "Table 2 (cont.)",
|
||||
"bbox": [100, 420, 600, 460],
|
||||
},
|
||||
]
|
||||
|
||||
inventory = build_table_inventory(structured_blocks)
|
||||
|
||||
assert len(inventory["tables"]) == 1
|
||||
t = inventory["tables"][0]
|
||||
assert t["is_continuation"] is True
|
||||
assert t["asset_block_id"] == "p5_a1"
|
||||
|
||||
|
||||
def test_multiple_captions_match_correct_assets_in_order() -> None:
|
||||
from paperforge.worker.ocr_tables import build_table_inventory
|
||||
|
||||
structured_blocks = [
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 5,
|
||||
"block_id": "a5b",
|
||||
"role": "table_asset",
|
||||
"text": "table 1 body",
|
||||
"bbox": [100, 100, 600, 400],
|
||||
},
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 5,
|
||||
"block_id": "c5a",
|
||||
"role": "table_caption",
|
||||
"text": "Table 1. First",
|
||||
"bbox": [100, 420, 600, 460],
|
||||
},
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 6,
|
||||
"block_id": "a6b",
|
||||
"role": "table_asset",
|
||||
"text": "table 2 body",
|
||||
"bbox": [100, 100, 600, 400],
|
||||
},
|
||||
{
|
||||
"paper_id": "KEY001",
|
||||
"page": 6,
|
||||
"block_id": "c6a",
|
||||
"role": "table_caption",
|
||||
"text": "Table 2. Second",
|
||||
"bbox": [100, 420, 600, 460],
|
||||
},
|
||||
]
|
||||
|
||||
inventory = build_table_inventory(structured_blocks)
|
||||
|
||||
assert inventory["official_table_count"] == 2
|
||||
assert inventory["tables"][0]["asset_block_id"] == "a5b"
|
||||
assert inventory["tables"][1]["asset_block_id"] == "a6b"
|
||||
|
|
|
|||
Loading…
Reference in a new issue