mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 06:50:53 +00:00
fix: stabilize OCR table matching and image-first rendering
This commit is contained in:
Research Assistant
parent
0e41454748
commit
174d2e8ccd
2 changed files with 315 additions and 38 deletions
|
|
@ -11,6 +11,8 @@ _TABLE_PREFIX_PATTERN = re.compile(
|
|||
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)
|
||||
|
|
@ -22,12 +24,31 @@ 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", "")
|
||||
|
|
@ -40,52 +61,90 @@ 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", "")
|
||||
caption_text = caption.get("text", "") or ""
|
||||
caption_bbox = caption.get("bbox", [0, 0, 0, 0])
|
||||
table_num = _extract_table_number(caption_text)
|
||||
is_continuation = bool(_CONTINUATION_PATTERN.search(caption_text))
|
||||
|
||||
candidate_pages = [caption_page, caption_page + 1, caption_page - 1]
|
||||
matched_asset = None
|
||||
# Continuations should only match same-page assets
|
||||
candidate_pages = [caption_page]
|
||||
if not is_continuation:
|
||||
candidate_pages.extend([caption_page + 1, caption_page - 1])
|
||||
|
||||
best_match: tuple[int | None, float] = (None, -999)
|
||||
for page in candidate_pages:
|
||||
if page < 1:
|
||||
continue
|
||||
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
|
||||
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)
|
||||
|
||||
tables.append(
|
||||
{
|
||||
"caption_block_id": caption.get("block_id", ""),
|
||||
"page": caption_page,
|
||||
"caption_text": caption_text,
|
||||
"table_number": table_num,
|
||||
"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 "",
|
||||
"truth_source": "image",
|
||||
"has_asset": matched_asset is not None,
|
||||
}
|
||||
)
|
||||
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])
|
||||
|
||||
for caption in captions:
|
||||
if not any(t["caption_block_id"] == caption.get("block_id", "") for t in tables if t["has_asset"]):
|
||||
if matched_asset is None:
|
||||
unmatched_captions.append(caption)
|
||||
|
||||
for i, asset in enumerate(assets):
|
||||
if i not in used_asset_indices:
|
||||
unmatched_assets.append(asset)
|
||||
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 "",
|
||||
"truth_source": "image",
|
||||
"has_asset": matched_asset is not None,
|
||||
})
|
||||
|
||||
unmatched_assets_list = [
|
||||
a for i, a in enumerate(assets) if i not in used_asset_indices
|
||||
]
|
||||
|
||||
official_count = len([t for t in tables if t["has_asset"] and not t["is_continuation"]])
|
||||
|
||||
return {
|
||||
"tables": tables,
|
||||
"unmatched_captions": unmatched_captions,
|
||||
"unmatched_assets": unmatched_assets,
|
||||
"official_table_count": len([t for t in tables if t["has_asset"]]),
|
||||
"unmatched_assets": unmatched_assets_list,
|
||||
"official_table_count": official_count,
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,8 +1,4 @@
|
|||
"""Phase 2 contract tests for table inventory.
|
||||
|
||||
paperforge.worker.ocr_tables does not exist yet -- tests will fail until
|
||||
Task 7 implements the module.
|
||||
"""
|
||||
"""Task 5 tests for hardened table matching."""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
|
|
@ -63,3 +59,225 @@ 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