From 53a5336b6af49ada432580ae0e24562dcbcc4f46 Mon Sep 17 00:00:00 2001
From: LLLin000 <809867916@qq.com>
Date: Sat, 4 Jul 2026 22:02:19 +0800
Subject: [PATCH] feat(pr3): add column-aware same-page reference boundary

---
 paperforge/worker/ocr_document.py | 76 ++++++++++++++++++++++++++++++-
 tests/test_ocr_document.py        | 47 ++++++++++++++++++-
 2 files changed, 121 insertions(+), 2 deletions(-)

diff --git a/paperforge/worker/ocr_document.py b/paperforge/worker/ocr_document.py
index c13948d5..0beacfe9 100644
--- a/paperforge/worker/ocr_document.py
+++ b/paperforge/worker/ocr_document.py
@@ -522,6 +522,52 @@ def _block_y_bottom(block: dict) -> float:
     return bbox[3] if bbox else 0.0
 
 
+
+def _block_column_band(block: dict, page_width: float) -> int | None:
+    """Classify a block into left (0), right (1), or center/mixed (None) column band.
+    Covers full-width blocks by checking width threshold.
+    """
+    bbox = block.get("bbox") or block.get("block_bbox") or [0, 0, 0, 0]
+    if len(bbox) < 4 or page_width <= 0:
+        return None
+    x0, x1 = bbox[0], bbox[2]
+    block_width = x1 - x0
+    # Full-width blocks span both columns → compatible with everything
+    if block_width >= page_width * 0.8:
+        return None  # full-width → treat as center/ambiguous
+    cx = (x0 + x1) / 2.0
+    if cx < page_width * 0.45:
+        return 0
+    if cx > page_width * 0.55:
+        return 1
+    return None
+
+
+def _is_full_width_ref_heading(block: dict, page_width: float) -> bool:
+    """True if the reference heading spans full page width (single-column layout)."""
+    bbox = block.get("bbox") or block.get("block_bbox") or [0, 0, 0, 0]
+    if len(bbox) < 4 or page_width <= 0:
+        return False
+    block_width = bbox[2] - bbox[0]
+    return block_width >= page_width * 0.8
+
+
+def _is_in_same_reference_column(block: dict, ref_heading_block: dict | None, page_width: float) -> bool:
+    """Check if block shares the same column band as the reference heading.
+    None ref_heading_block -> True (conservative).
+    Full-width ref heading -> True (page-level, not column-level).
+    """
+    if ref_heading_block is None:
+        return True
+    ref_band = _block_column_band(ref_heading_block, page_width)
+    if ref_band is None:
+        return True
+    block_band = _block_column_band(block, page_width)
+    if block_band is None:
+        return True
+    return block_band == ref_band
+
+
 _REFERENCE_ZONE_MARKER_TYPES: frozenset[str] = frozenset({
     "reference_numeric_bracket",
     "reference_numeric_dot",
@@ -1246,6 +1292,23 @@ def _is_below_same_page_reference_heading(block: dict, refs_start_page: int | No
     return _block_y_top(block) > ref_heading_top
 
 
+
+
+def _page_width_for_zone_block(block: dict, page_blocks: list[dict], body_anchor: dict | None = None) -> float:
+    """Derive page_width from block, page_blocks, or body_anchor. Fallback 1200."""
+    pw = block.get("page_width", 0) or 0
+    if pw:
+        return float(pw)
+    for b in page_blocks:
+        bw = b.get("page_width", 0) or 0
+        if bw:
+            pw = max(pw, float(bw))
+    if pw:
+        return pw
+    if body_anchor and isinstance(body_anchor, dict):
+        pw = body_anchor.get("page_width", 0) or 0
+    return float(pw) if pw else 1200.0
+
 def infer_zones(
     blocks: list[dict],
     anchors: dict[str, dict] | None,
@@ -1524,6 +1587,13 @@ def infer_zones(
                 ref_heading_top = _block_y_top(block)
                 break
 
+    # Derive page_width for column-aware reference boundary checks
+    _zone_page_width = _page_width_for_zone_block(
+        ref_heading_block if ref_heading_block else (blocks[0] if blocks else {}),
+        blocks,
+        anchors,
+    )
+
     body_blocks = [
         block
         for block in blocks
@@ -1539,7 +1609,10 @@ def infer_zones(
         )
         and (
             (body_end_page is None or int(block.get("page", 0) or 0) <= body_end_page)
-            or _is_above_same_page_reference_heading(block, refs_start_page, ref_heading_top if ref_heading_block else None)
+            or (
+                _is_above_same_page_reference_heading(block, refs_start_page, ref_heading_top if ref_heading_block else None)
+                and _is_in_same_reference_column(block, ref_heading_block, _zone_page_width)
+            )
         )
         and not _is_reference_item_candidate(block)
         and not _is_reference_heading_candidate(block)
@@ -1555,6 +1628,7 @@ def infer_zones(
         block
         for block in blocks
         if _is_below_same_page_reference_heading(block, refs_start_page, ref_heading_top if ref_heading_block else None)
+        and _is_in_same_reference_column(block, ref_heading_block, _zone_page_width)
         and not _is_reference_item_candidate(block)
         and not _is_reference_heading_candidate(block)
         and block.get("block_id") is not None
diff --git a/tests/test_ocr_document.py b/tests/test_ocr_document.py
index d2ae4258..2b24d386 100644
--- a/tests/test_ocr_document.py
+++ b/tests/test_ocr_document.py
@@ -5555,4 +5555,49 @@ def test_frontmatter_support_below_body_start_rescued_by_width() -> None:
     assert "corresp" in fm_main["block_ids"], \
         "narrow frontmatter_support block should be in frontmatter_main_zone"
     assert "corresp" not in body["block_ids"], \
-        "narrow frontmatter_support block should NOT be in body_zone"
\ No newline at end of file
+        "narrow frontmatter_support block should NOT be in body_zone"
+
+def test_block_column_band_left_right() -> None:
+    from paperforge.worker.ocr_document import _block_column_band
+    pw = 1200.0
+    left = {"bbox": [100, 100, 500, 200]}
+    right = {"bbox": [700, 100, 1100, 200]}
+    center = {"bbox": [200, 100, 1000, 200]}
+    full = {"bbox": [50, 100, 1150, 200]}
+    assert _block_column_band(left, pw) == 0
+    assert _block_column_band(right, pw) == 1
+    assert _block_column_band(center, pw) is None
+    assert _block_column_band(full, pw) is None  # full-width -> None
+
+
+def test_is_in_same_reference_column() -> None:
+    from paperforge.worker.ocr_document import _is_in_same_reference_column
+    pw = 1200.0
+    ref_heading = {"bbox": [700, 100, 1100, 150]}  # right column
+    left_block = {"bbox": [100, 400, 500, 500]}    # left column
+    right_block = {"bbox": [700, 400, 1100, 500]}  # right column
+    assert not _is_in_same_reference_column(left_block, ref_heading, pw)
+    assert _is_in_same_reference_column(right_block, ref_heading, pw)
+    assert _is_in_same_reference_column(left_block, None, pw)  # None ref -> True
+
+
+def test_is_in_same_reference_column_full_width_ref() -> None:
+    from paperforge.worker.ocr_document import _is_in_same_reference_column
+    pw = 1200.0
+    full_ref = {"bbox": [50, 100, 1150, 150]}  # full-width
+    left_block = {"bbox": [100, 400, 500, 500]}
+    assert _is_in_same_reference_column(left_block, full_ref, pw)  # full-width -> page-level
+
+
+def test_same_page_tail_column_aware() -> None:
+    from paperforge.worker.ocr_document import infer_zones
+    blocks = [
+        {"page": 3, "role": "reference_heading", "text": "REFERENCES", "bbox": [700, 100, 1100, 140], "block_id": "p3_h1"},
+        {"page": 3, "role": "reference_item", "text": "[1] Some ref", "bbox": [700, 150, 1100, 180], "block_id": "p3_r1"},
+        {"page": 3, "role": "body_paragraph", "text": "Conclusions continue...", "bbox": [100, 300, 500, 350], "block_id": "p3_b1"},
+    ]
+    anchors = {"body_family_anchor": {"status": "ACCEPT", "sample_pages": [2]}}
+    region_bus = infer_zones(blocks, anchors)
+    # Left-column body below ref heading but in different column -> NOT in tail zone
+    assert "p3_b1" not in region_bus["tail_nonref_hold_zone"]["block_ids"], \
+        "Left-column body should not be in tail_nonref_hold_zone despite being below ref heading"