diff --git a/paperforge/worker/ocr_figure_reader.py b/paperforge/worker/ocr_figure_reader.py new file mode 100644 index 00000000..e6fa930f --- /dev/null +++ b/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, + } diff --git a/tests/test_ocr_figure_reader.py b/tests/test_ocr_figure_reader.py new file mode 100644 index 00000000..6c781a3b --- /dev/null +++ b/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]