lllin000_PaperForge/paperforge/worker/ocr_figure_reader.py
LLLin000 1918f392db fix: address 5 post-review issues in previous-page locator bridge
1. cluster_bbox: compute from consumed asset bboxes, add to matched entry
2. Visual group selection: prioritize candidate_groups (composite_parent,
   distance_cluster) before falling back to loose asset collection.
   Fallback only fires for tight single clusters (_is_tight_asset_cluster).
3. Page-aware key: unmatched_legends and ambiguous_figures cleanup now
   uses (page, block_id) tuple, not bare block_id.
4. Reader bridge consumption: _normalize_bucket passes bridge_block_ids
   through to normalized output; _materialize_reader_figure includes
   bridge blocks in consumed_caption_block_ids.
5. Stale _unowned_by_page: visual group search uses used_asset_page_ids
   to avoid double-consumption across multiple locators on one page.
2026-06-29 18:37:26 +08:00

573 lines
22 KiB
Python

from __future__ import annotations
import contextlib
import re
from dataclasses import dataclass
_FIGURE_NUM_PATTERN = re.compile(
r"(?:Figure|Fig\.?|Supplementary\s+Figure|Supplementary\s+Fig\.?|"
r"Extended\s+Data\s+Figure|Extended\s+Data\s+Fig\.?)\s+"
r"(?:S)?(\d+(?:\.\d+)?)",
flags=re.IGNORECASE,
)
@dataclass(frozen=True)
class ReaderCoverage:
total: int
accounted: int
gap_count: int
def as_dict(self) -> dict:
ratio = 1.0 if self.total == 0 else self.accounted / self.total
return {
"total": self.total,
"accounted": self.accounted,
"gap_count": self.gap_count,
"ratio": ratio,
}
def _index_structured_blocks(structured_blocks: list[dict]) -> dict[int | str, dict]:
return {block.get("block_id"): block for block in structured_blocks if block.get("block_id") is not None}
def _index_blocks_by_page_and_id(blocks: list[dict]) -> dict[tuple[int | str | None, int | str], dict]:
return {(block.get("page"), block.get("block_id")): block for block in blocks if block.get("block_id") is not None}
def _legend_block_id(item: dict) -> int | str | None:
return item.get("legend_block_id", item.get("block_id"))
def _caption_text(item: dict) -> str:
return str(item.get("caption_text", item.get("text", "")) or "")
def _asset_ids_from_item(item: dict) -> list[int | str]:
if "asset_block_ids" in item:
return list(item.get("asset_block_ids", []))
if "matched_assets" in item:
return [asset.get("block_id") for asset in item.get("matched_assets", []) if asset.get("block_id") is not None]
if "candidates" in item:
return [
asset.get("asset_block_id")
for asset in item.get("candidates", [])
if asset.get("asset_block_id") is not None
]
if "media_block_ids" in item:
return list(item.get("media_block_ids", []))
return []
def _candidate_asset_ids_from_item(item: dict) -> list[int | str]:
if "candidate_asset_ids" in item:
return list(item.get("candidate_asset_ids", []))
if "candidates" in item:
return [
asset.get("asset_block_id")
for asset in item.get("candidates", [])
if asset.get("asset_block_id") is not None
]
return []
def _formal_legend_blocks(blocks: list[dict]) -> list[dict]:
"""Return blocks that are formal figure legends, excluding table_caption."""
legends = []
for block in blocks:
role = block.get("role")
if role == "table_caption":
continue
marker_type = str((block.get("marker_signature") or {}).get("type") or "")
style_family = str(block.get("style_family") or "")
if role in {"figure_caption", "legend"}:
legends.append(block)
continue
if marker_type == "figure_number" and style_family in {"legend_like"}:
legends.append(block)
return legends
def _normalize_bucket(
items: list[dict],
bucket_name: str,
block_index: dict[int | str, dict],
legends_index: dict[int | str, dict] | None = None,
*,
page_block_index: dict | None = None,
page_legends_index: dict | None = None,
) -> list[dict]:
candidate_buckets = {"held_figures", "ambiguous_figures"}
normalized = []
for item in items:
source_item = item
legend_block_id = _legend_block_id(source_item)
source_page = source_item.get("page")
block = (page_block_index or {}).get((source_page, legend_block_id), block_index.get(legend_block_id, {}))
source_text = str(source_item.get("text", source_item.get("caption_text", block.get("text", ""))) or "")
legend_data = (page_legends_index or {}).get(
(source_page, legend_block_id), (legends_index or {}).get(legend_block_id, {})
)
source_marker = (
source_item.get("marker_signature")
or legend_data.get("marker_signature")
or block.get("marker_signature")
or {}
)
_inferred_figure_number: int | None = None
_inferred_marker_type: str | None = None
_text_for_inference = source_text or block.get("text", "")
if not source_item.get("marker_type") and not source_marker.get("type") and _text_for_inference:
_m = _FIGURE_NUM_PATTERN.search(_text_for_inference)
if _m:
_inferred_marker_type = "figure_number"
with contextlib.suppress(ValueError):
_inferred_figure_number = int(float(_m.group(1)))
normalized.append(
{
"figure_id": source_item.get("figure_id", ""),
"figure_number": source_item.get("figure_number")
or source_marker.get("number")
or _inferred_figure_number,
"legend_block_id": legend_block_id,
"caption_text": source_text,
"asset_block_ids": _asset_ids_from_item(source_item),
"candidate_asset_ids": (
_candidate_asset_ids_from_item(source_item) if bucket_name in candidate_buckets else []
),
"marker_type": source_item.get("marker_type") or source_marker.get("type") or _inferred_marker_type,
"inline_mention": bool(source_item.get("inline_mention", False)),
"panel_label": bool(source_item.get("panel_label", False)),
"body_prose_likelihood": float(source_item.get("body_prose_likelihood", 0.0)),
"strict_reject": bool(source_item.get("strict_reject", False)),
"linked_legend_block_id": source_item.get("linked_legend_block_id"),
"cluster_area_ratio": float(source_item.get("cluster_area_ratio", 0.0)),
"width_ratio": float(source_item.get("width_ratio", 0.0)),
"height_ratio": float(source_item.get("height_ratio", 0.0)),
"media_block_count": int(source_item.get("media_block_count", len(_asset_ids_from_item(source_item)))),
"page": source_item.get("page", block.get("page")),
"legend_page": source_item.get("legend_page", source_page),
"zone": (legend_data.get("zone") or source_item.get("zone") or block.get("zone")),
"style_family": (
legend_data.get("style_family") or source_item.get("style_family") or block.get("style_family")
),
"strict_status": source_item.get("strict_status", bucket_name.removesuffix("s")),
"bridge_block_ids": source_item.get("bridge_block_ids", []),
"source_item": source_item,
}
)
return normalized
def _build_bucket(
strict_inventory: dict,
key: str,
block_index: dict[int | str, dict],
legends_index: dict[int | str, dict] | None = None,
*,
page_block_index: dict | None = None,
page_legends_index: dict | None = None,
) -> list[dict]:
return _normalize_bucket(
strict_inventory.get(key, []),
key,
block_index,
legends_index,
page_block_index=page_block_index,
page_legends_index=page_legends_index,
)
def _caption_consumption_page(item: dict) -> int | str | None:
return item.get("legend_page", item.get("page"))
_READER_STATUS_MAP = {
"matched_figures": "EXACT_MATCH",
"held_figures": "HOLD",
"ambiguous_figures": "GROUPED_APPROXIMATE",
"unmatched_legends": "LEGEND_ONLY",
"unresolved_clusters": "ASSET_GROUP_ONLY",
}
def _stable_reader_figure_id(
figure_number: int | None,
*,
page: int | None = None,
first_asset_block_id: int | str | None = None,
ordinal: int | None = None,
suffix: str = "",
) -> str:
if figure_number is not None:
base = f"figure_{int(figure_number):03d}_reader"
elif first_asset_block_id is not None:
base = f"visual_group_{int(page or 0)}_{first_asset_block_id}_reader"
else:
base = f"visual_group_{int(page or 0)}_{int(ordinal or 0)}_reader"
return f"{base}{suffix}"
def _materialize_hold_outcome(
*,
legend_block_id: int | str | None,
caption_text: str,
page: int | None,
candidate_asset_ids: list[int | str],
hold_visibility: str,
) -> dict:
reader_visible = hold_visibility == "reader_hold"
return {
"reader_figure_id": _stable_reader_figure_id(None, page=page, ordinal=0, suffix="_hold"),
"figure_number": None,
"reader_status": "HOLD",
"strict_status": "held",
"strict_source": "held_figures",
"caption_block_id": legend_block_id,
"caption_text": caption_text,
"visual_groups": [],
"consumed_caption_block_ids": (
[{"page": page, "block_id": legend_block_id}]
if reader_visible and legend_block_id is not None and caption_text
else []
),
"consumed_asset_block_ids": [],
"debug_refs": {"candidate_asset_ids": list(candidate_asset_ids), "hold_visibility": hold_visibility},
}
def _placeholder_visual_group(
*,
page: int | None,
status: str,
candidate_asset_ids: list[int | str] | None = None,
rendered_as_representative: bool = False,
) -> dict:
return {
"page": page,
"asset_block_ids": list(candidate_asset_ids or []),
"group_status": status,
"rendered_as_representative": rendered_as_representative,
}
def _materialize_reader_figure(
normalized_item: dict,
source_bucket: str,
ordinal: int,
) -> dict | None:
figure_number = normalized_item.get("figure_number")
caption_text = normalized_item.get("caption_text", "")
legend_block_id = normalized_item.get("legend_block_id")
if source_bucket == "matched_figures":
asset_ids = list(normalized_item.get("asset_block_ids", []))
strict_status = normalized_item.get("strict_status", "matched")
reader_status = "SEQUENCE_MATCH" if strict_status == "sequence_match" else "EXACT_MATCH"
return {
"reader_figure_id": _stable_reader_figure_id(
figure_number,
page=normalized_item.get("page"),
first_asset_block_id=asset_ids[0] if asset_ids else None,
ordinal=ordinal,
),
"figure_id": normalized_item.get("figure_id", ""),
"figure_number": figure_number,
"reader_status": reader_status,
"strict_status": strict_status,
"strict_source": source_bucket,
"caption_block_id": legend_block_id,
"caption_text": caption_text,
"visual_groups": [
{
"page": normalized_item.get("page"),
"asset_block_ids": asset_ids,
"group_status": "matched_group",
"rendered_as_representative": True,
}
],
"consumed_caption_block_ids": (([{"page": _caption_consumption_page(normalized_item), "block_id": legend_block_id}] if legend_block_id is not None else [])
+ [{"page": normalized_item.get("page"), "block_id": bid} for bid in normalized_item.get("bridge_block_ids", [])]),
"consumed_asset_block_ids": [{"page": normalized_item.get("page"), "block_id": aid} for aid in asset_ids],
"debug_refs": {},
}
if source_bucket in {"held_figures", "ambiguous_figures"}:
candidate_asset_ids = list(normalized_item.get("candidate_asset_ids", []))
if candidate_asset_ids:
return {
"reader_figure_id": _stable_reader_figure_id(
figure_number,
page=normalized_item.get("page"),
first_asset_block_id=candidate_asset_ids[0] if candidate_asset_ids else None,
ordinal=ordinal,
),
"figure_number": figure_number,
"reader_status": "GROUPED_APPROXIMATE",
"strict_status": normalized_item.get("strict_status", "ambiguous"),
"strict_source": source_bucket,
"caption_block_id": legend_block_id,
"caption_text": caption_text,
"visual_groups": [
{
"page": normalized_item.get("page"),
"asset_block_ids": candidate_asset_ids,
"group_status": "candidate_group",
"rendered_as_representative": False,
}
],
"consumed_caption_block_ids": [{"page": _caption_consumption_page(normalized_item), "block_id": legend_block_id}]
if legend_block_id is not None
else [],
"consumed_asset_block_ids": [],
"debug_refs": {"candidate_asset_ids": candidate_asset_ids},
}
if caption_text:
return {
"reader_figure_id": _stable_reader_figure_id(
figure_number,
page=normalized_item.get("page"),
ordinal=ordinal,
),
"figure_number": figure_number,
"reader_status": "LEGEND_ONLY",
"strict_status": normalized_item.get("strict_status", "held"),
"strict_source": source_bucket,
"caption_block_id": legend_block_id,
"caption_text": caption_text,
"visual_groups": [
_placeholder_visual_group(
page=normalized_item.get("page"),
status="legend_only_group",
)
],
"consumed_caption_block_ids": [{"page": _caption_consumption_page(normalized_item), "block_id": legend_block_id}]
if legend_block_id is not None
else [],
"consumed_asset_block_ids": [],
"debug_refs": {},
}
return None
if source_bucket == "unmatched_legends":
return {
"reader_figure_id": _stable_reader_figure_id(
figure_number,
page=normalized_item.get("page"),
ordinal=ordinal,
),
"figure_number": figure_number,
"reader_status": "LEGEND_ONLY",
"strict_status": normalized_item.get("strict_status", "unmatched"),
"strict_source": source_bucket,
"caption_block_id": legend_block_id,
"caption_text": caption_text,
"visual_groups": [
_placeholder_visual_group(
page=normalized_item.get("page"),
status="legend_only_group",
)
],
"consumed_caption_block_ids": [{"page": _caption_consumption_page(normalized_item), "block_id": legend_block_id}]
if legend_block_id is not None
else [],
"consumed_asset_block_ids": [],
"debug_refs": {},
}
if source_bucket == "unresolved_clusters":
asset_ids = list(normalized_item.get("asset_block_ids", []))
return {
"reader_figure_id": _stable_reader_figure_id(
None,
page=normalized_item.get("page"),
first_asset_block_id=asset_ids[0] if asset_ids else None,
ordinal=ordinal,
),
"figure_number": None,
"reader_status": "ASSET_GROUP_ONLY",
"strict_status": normalized_item.get("strict_status", "unresolved_cluster"),
"strict_source": source_bucket,
"caption_block_id": None,
"caption_text": "",
"visual_groups": [
{
"page": normalized_item.get("page"),
"asset_block_ids": asset_ids,
"group_status": "candidate_group",
"rendered_as_representative": True,
}
],
"consumed_caption_block_ids": [],
"consumed_asset_block_ids": [{"page": normalized_item.get("page"), "block_id": aid} for aid in asset_ids],
"debug_refs": {},
}
return None
def _passes_formal_legend_gate(item: dict) -> bool:
return (
item.get("marker_type") == "figure_number"
and not bool(item.get("inline_mention"))
and not bool(item.get("panel_label"))
and float(item.get("body_prose_likelihood", 0.0)) < 0.5
and not bool(item.get("strict_reject"))
)
def _passes_salient_visual_group_gate(item: dict) -> bool:
zone = str(item.get("zone") or "")
if zone == "preproof_cover_zone":
return False
area_ratio = float(item.get("cluster_area_ratio", 0.0))
width_ratio = float(item.get("width_ratio", 0.0))
height_ratio = float(item.get("height_ratio", 0.0))
media_count = int(item.get("media_block_count", len(item.get("asset_block_ids", []))))
if area_ratio >= 0.03:
return True
if width_ratio >= 0.30 and height_ratio >= 0.08:
return True
if media_count >= 3:
return True
return bool(media_count >= 2 and area_ratio >= 0.02)
def _collect_reader_eligible_inputs(normalized: dict) -> list[dict]:
eligible: list[dict] = []
seen_legends: set[tuple[int | str | None, int | str | None]] = set()
for source_name in ("matched_figures", "held_figures", "ambiguous_figures"):
for item in normalized.get(source_name, []):
legend_id = (item.get("page"), item.get("legend_block_id"))
if legend_id in seen_legends:
continue
if not _passes_formal_legend_gate(item):
continue
seen_legends.add(legend_id)
eligible.append({"kind": "legend", "source": source_name, "item": item})
for item in normalized.get("unmatched_legends", []):
legend_id = (item.get("page"), item.get("legend_block_id"))
if legend_id in seen_legends:
continue
if not _passes_formal_legend_gate(item):
continue
seen_legends.add(legend_id)
eligible.append({"kind": "legend", "source": "unmatched_legends", "item": item})
for item in normalized.get("unresolved_clusters", []):
if item.get("linked_legend_block_id") is not None:
continue
if not _passes_salient_visual_group_gate(item):
continue
eligible.append({"kind": "visual_group", "source": "unresolved_clusters", "item": item})
return eligible
def synthesize_reader_figures(
strict_inventory: dict,
structured_blocks: list[dict],
document_structure: object | None = None,
) -> dict:
normalized = _normalize_strict_figure_inventory(strict_inventory, structured_blocks)
eligible_inputs = _collect_reader_eligible_inputs(normalized)
reader_figures: list[dict] = []
consumed_caption_ids: list[int | str] = []
consumed_asset_ids: list[int | str] = []
for ordinal, entry in enumerate(eligible_inputs):
item = entry["item"]
source_name = entry["source"]
materialized = _materialize_reader_figure(item, source_name, ordinal)
if materialized is None:
continue
reader_figures.append(materialized)
consumed_caption_ids.extend(materialized.get("consumed_caption_block_ids", []))
consumed_asset_ids.extend(materialized.get("consumed_asset_block_ids", []))
deduped_legend_ids: list[dict] = strict_inventory.get("deduped_legend_ids", [])
for item in deduped_legend_ids:
page = item.get("page") if isinstance(item, dict) else None
block_id = item.get("block_id") if isinstance(item, dict) else item
if block_id is not None:
consumed_caption_ids.append({"page": page, "block_id": block_id})
coverage_total = len(eligible_inputs)
return {
"normalized_inputs": normalized,
"reader_figures": reader_figures,
"reader_coverage": ReaderCoverage(
total=coverage_total,
accounted=len(reader_figures),
gap_count=coverage_total - len(reader_figures),
).as_dict(),
"consumed_caption_block_ids": consumed_caption_ids,
"consumed_asset_block_ids": consumed_asset_ids,
}
def _normalize_strict_figure_inventory(
strict_inventory: dict,
structured_blocks: list[dict],
) -> dict:
block_index = _index_structured_blocks(structured_blocks)
page_block_index = _index_blocks_by_page_and_id(structured_blocks)
legends_index: dict[int | str, dict] = {
item.get("block_id"): item
for item in strict_inventory.get("figure_legends", [])
if item.get("block_id") is not None
}
page_legends_index: dict[tuple[int | str | None, int | str], dict] = {
(item.get("page"), item.get("block_id")): item
for item in strict_inventory.get("figure_legends", [])
if item.get("block_id") is not None
}
return {
"matched_figures": _build_bucket(
strict_inventory,
"matched_figures",
block_index,
legends_index,
page_block_index=page_block_index,
page_legends_index=page_legends_index,
),
"held_figures": _build_bucket(
strict_inventory,
"held_figures",
block_index,
legends_index,
page_block_index=page_block_index,
page_legends_index=page_legends_index,
),
"ambiguous_figures": _build_bucket(
strict_inventory,
"ambiguous_figures",
block_index,
legends_index,
page_block_index=page_block_index,
page_legends_index=page_legends_index,
),
"unmatched_legends": _build_bucket(
strict_inventory,
"unmatched_legends",
block_index,
legends_index,
page_block_index=page_block_index,
page_legends_index=page_legends_index,
),
"unresolved_clusters": _build_bucket(
strict_inventory,
"unresolved_clusters",
block_index,
legends_index,
page_block_index=page_block_index,
page_legends_index=page_legends_index,
),
"block_index": block_index,
}