feat: normalize strict OCR figure inventory for reader synthesis

paperforge/worker/ocr_figure_reader.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@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 _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 _normalize_bucket(
+    items: list[dict],
+    bucket_name: str,
+    block_index: dict[int | str, dict],
+) -> list[dict]:
+    candidate_buckets = {"held_figures", "ambiguous_figures"}
+    normalized = []
+    for item in items:
+        legend_block_id = _legend_block_id(item)
+        block = block_index.get(legend_block_id, {})
+        normalized.append(
+            {
+                "figure_number": item.get("figure_number"),
+                "legend_block_id": legend_block_id,
+                "caption_text": _caption_text(item),
+                "asset_block_ids": _asset_ids_from_item(item),
+                "candidate_asset_ids": (_asset_ids_from_item(item) if bucket_name in candidate_buckets else []),
+                "marker_type": item.get("marker_type") or (block.get("marker_signature") or {}).get("type"),
+                "inline_mention": bool(item.get("inline_mention", False)),
+                "panel_label": bool(item.get("panel_label", False)),
+                "body_prose_likelihood": float(item.get("body_prose_likelihood", 0.0)),
+                "zone": item.get("zone") or block.get("zone"),
+                "style_family": item.get("style_family") or block.get("style_family"),
+                "strict_status": item.get("strict_status", bucket_name.removesuffix("s")),
+                "source_item": item,
+            }
+        )
+    return normalized
+
+
+def _build_bucket(
+    strict_inventory: dict,
+    key: str,
+    block_index: dict[int | str, dict],
+) -> list[dict]:
+    return _normalize_bucket(strict_inventory.get(key, []), key, block_index)
+
+
+def _normalize_strict_figure_inventory(
+    strict_inventory: dict,
+    structured_blocks: list[dict],
+) -> dict:
+    block_index = _index_structured_blocks(structured_blocks)
+    return {
+        "matched_figures": _build_bucket(strict_inventory, "matched_figures", block_index),
+        "held_figures": _build_bucket(strict_inventory, "held_figures", block_index),
+        "ambiguous_figures": _build_bucket(strict_inventory, "ambiguous_figures", block_index),
+        "unmatched_legends": _build_bucket(strict_inventory, "unmatched_legends", block_index),
+        "unresolved_clusters": _build_bucket(strict_inventory, "unresolved_clusters", block_index),
+        "block_index": block_index,
+    }

tests/test_ocr_figure_reader.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+
+def test_normalize_strict_inventory_maps_bucket_variants_to_common_fields() -> None:
+    from paperforge.worker.ocr_figure_reader import _normalize_strict_figure_inventory
+
+    strict_inventory = {
+        "matched_figures": [
+            {
+                "figure_number": 6,
+                "block_id": 15,
+                "text": "Fig. 6 The figure represents...",
+                "matched_assets": [{"block_id": 40, "bbox": [1, 2, 3, 4]}],
+                "match_score": 0.91,
+            }
+        ],
+        "ambiguous_figures": [
+            {
+                "figure_number": 3,
+                "legend_block_id": 9,
+                "text": "FIGURE 3 | Histological evaluation...",
+                "candidates": [{"asset_block_id": 10, "match_score": 0.51}, {"asset_block_id": 11, "match_score": 0.49}],
+            }
+        ],
+        "unmatched_legends": [
+            {
+                "block_id": 21,
+                "text": "FIGURE 2 | Treadmill exercise protocols...",
+                "figure_number": 2,
+            }
+        ],
+        "unresolved_clusters": [
+            {
+                "page": 7,
+                "media_block_ids": [30, 31],
+            }
+        ],
+    }
+
+    structured_blocks = [
+        {"block_id": 21, "marker_signature": {"type": "figure_number"}, "zone": "display_zone", "style_family": "legend_like"}
+    ]
+
+    normalized = _normalize_strict_figure_inventory(strict_inventory, structured_blocks)
+
+    assert normalized["matched_figures"][0]["legend_block_id"] == 15
+    assert normalized["matched_figures"][0]["caption_text"] == "Fig. 6 The figure represents..."
+    assert normalized["matched_figures"][0]["asset_block_ids"] == [40]
+    assert normalized["ambiguous_figures"][0]["candidate_asset_ids"] == [10, 11]
+    assert normalized["unmatched_legends"][0]["legend_block_id"] == 21
+    assert normalized["unresolved_clusters"][0]["asset_block_ids"] == [30, 31]