"""Post-extraction recovery pass for orphan figures and tables.
Runs AFTER initial figure/table extraction, BEFORE completeness grading.
Audits results, detects orphans and floating captions, attempts targeted
pairing via y-proximity matching and numbering-gap search.
Does NOT create new figure/table objects — only fills in missing captions
on existing ones.
"""
from __future__ import annotations
import logging
import re
import pymupdf
from .models import ExtractedFigure, ExtractedTable, SectionSpan, PageExtraction
logger = logging.getLogger(__name__)
# Caption number extraction — matches plain digits or appendix-style (A.1, S1)
_CAPTION_NUM_RE = re.compile(r"(\d+|[A-Z]\.\d+|S\d+)")
# Relaxed caption patterns for gap search (not block-start anchored)
_NUM_GROUP = r"(\d+|[IVXLCDM]+|[A-Z]\.\d+|S\d+)"
_FIG_REF_RE = re.compile(
rf"(?:Figure|Fig\.?)\s+{_NUM_GROUP}\s*[.:()\u2014\u2013-]",
re.IGNORECASE,
)
_TABLE_REF_RE = re.compile(
rf"(?:Table|Tab\.)\s+{_NUM_GROUP}\s*[.:()\u2014\u2013-]",
re.IGNORECASE,
)
# Body-text reference context (reject mid-paragraph references)
_BODY_REF_CONTEXT_RE = re.compile(
r"(?:as\s+(?:shown|seen|described|illustrated|presented|depicted|listed|"
r"reported|summarized|given|indicated)\s+in\s+|"
r"(?:see|refer\s+to|shown\s+in|cf\.?)\s+)",
re.IGNORECASE,
)
def _adaptive_max_y_distance(
doc: pymupdf.Document,
*,
page_fraction: float = 0.15,
) -> float:
"""Compute fallback y-proximity threshold from page height.
Used when fewer than 3 matched caption-item pairs are available
for calibration. Returns *page_fraction* of the median page height.
"""
_ASSUMED_PAGE_HEIGHT = 792.0 # letter-size fallback
if not doc or len(doc) == 0:
return _ASSUMED_PAGE_HEIGHT * page_fraction
heights = [doc[i].rect.height for i in range(min(len(doc), 5))]
heights.sort()
median_height = heights[len(heights) // 2]
return median_height * page_fraction
def run_recovery(
doc: pymupdf.Document,
figures: list[ExtractedFigure],
tables: list[ExtractedTable],
page_chunks: list[dict],
sections: list[SectionSpan] | None = None,
pages: list[PageExtraction] | None = None,
) -> tuple[list[ExtractedFigure], list[ExtractedTable]]:
"""Run post-extraction recovery to fill orphan captions.
Args:
doc: Open pymupdf.Document.
figures: Extracted figures (may have None captions).
tables: Extracted tables (may have None captions).
page_chunks: Output of pymupdf4llm.to_markdown(page_chunks=True).
sections: Section spans (unused currently, reserved for future).
pages: Page extractions (unused currently, reserved for future).
Returns:
Updated (figures, tables) with recovered captions filled in.
"""
from .feature_extraction.captions import find_all_captions
max_y_dist = _adaptive_max_y_distance(doc)
fig_recoveries = _recover_captions(
doc, figures, find_all_captions,
_FIG_REF_RE, kind="figure", max_y_distance=max_y_dist,
)
tab_recoveries = _recover_captions(
doc, tables, find_all_captions,
_TABLE_REF_RE, kind="table", max_y_distance=max_y_dist,
)
if fig_recoveries or tab_recoveries:
logger.info(
"Recovery pass: %d figure caption(s), %d table caption(s) recovered",
fig_recoveries, tab_recoveries,
)
return figures, tables
def _recover_captions(
doc: pymupdf.Document,
items: list[ExtractedFigure] | list[ExtractedTable],
caption_finder,
ref_re: re.Pattern,
kind: str,
max_y_distance: float,
) -> int:
"""Core recovery logic for either figures or tables.
Returns the number of captions recovered.
"""
if not items:
return 0
include_figures = kind == "figure"
include_tables = kind == "table"
# --- Step 1: Audit ---
# Map caption numbers (as strings) to item indices
assigned_nums: dict[str, int] = {}
orphan_indices: list[int] = []
for idx, item in enumerate(items):
if item.caption:
m = _CAPTION_NUM_RE.search(item.caption)
if m:
assigned_nums[m.group(1)] = idx
else:
orphan_indices.append(idx)
if not orphan_indices:
return 0 # nothing to recover
# Scan all pages for captions -> find floating (unassigned) ones
# Also collect distances from assigned captions to their items for calibration
floating_captions: list[dict] = [] # {num, y_center, text, page_num}
matched_distances: list[float] = []
for page_num_0, page in enumerate(doc):
page_num = page_num_0 + 1
detected = caption_finder(
page,
include_figures=include_figures,
include_tables=include_tables,
)
for cap in detected:
num_str = cap.number
if not num_str:
continue
if num_str in assigned_nums:
# Measure distance from this caption to its matched item
item_idx = assigned_nums[num_str]
item = items[item_idx]
if item.page_num == page_num:
dist = abs(cap.y_center - item.bbox[3])
matched_distances.append(dist)
else:
floating_captions.append({
"num": num_str,
"y_center": cap.y_center,
"text": cap.text,
"page_num": page_num,
"bbox": cap.bbox,
})
# Calibrate threshold from matched caption-to-item distances
if len(matched_distances) >= 3:
matched_distances.sort()
n_md = len(matched_distances)
q1 = matched_distances[n_md // 4]
q3 = matched_distances[3 * n_md // 4]
iqr = q3 - q1
effective_max_dist = q3 + 1.5 * iqr
logger.debug(
"Recovery %s: calibrated max_y_distance=%.0f from %d matched pairs "
"(Q1=%.0f Q3=%.0f IQR=%.0f, passed default=%.0f)",
kind, effective_max_dist, n_md, q1, q3, iqr, max_y_distance,
)
else:
effective_max_dist = max_y_distance
# --- Step 2: Y-Proximity Matching (orphans ↔ floating captions) ---
recoveries = 0
matched_floating: set[int] = set() # indices into floating_captions
# Sort orphans top-to-bottom within each page for greedy assignment
orphan_indices_sorted = sorted(
orphan_indices,
key=lambda i: (items[i].page_num, items[i].bbox[1]),
)
for orphan_idx in orphan_indices_sorted:
orphan = items[orphan_idx]
orphan_page = orphan.page_num
orphan_bottom = orphan.bbox[3] # y1
best_dist = effective_max_dist + 1
best_fc_idx = -1
for fc_idx, fc in enumerate(floating_captions):
if fc_idx in matched_floating:
continue
if fc["page_num"] != orphan_page:
continue
dist = abs(fc["y_center"] - orphan_bottom)
if dist < best_dist:
best_dist = dist
best_fc_idx = fc_idx
if best_fc_idx >= 0 and best_dist <= effective_max_dist:
items[orphan_idx].caption = floating_captions[best_fc_idx]["text"]
num = floating_captions[best_fc_idx]["num"]
assigned_nums[num] = orphan_idx
matched_floating.add(best_fc_idx)
recoveries += 1
logger.debug(
"Recovery: %s orphan p%d matched to floating caption %s (dist=%.0f)",
kind, orphan_page, num, best_dist,
)
# Refresh orphan list after proximity matching
orphan_indices = [i for i in orphan_indices if items[i].caption is None]
if not orphan_indices:
return recoveries
# --- Step 3: Numbering Gap Search ---
# Only works with integer-numbered captions (not appendix "A.1", "S1" etc.)
if not assigned_nums:
return recoveries
int_nums: dict[int, int] = {}
for k, v in assigned_nums.items():
try:
int_nums[int(k)] = v
except ValueError:
pass # skip appendix-style numbers
if not int_nums:
return recoveries
min_num = min(int_nums)
max_num = max(int_nums)
gap_nums = [n for n in range(min_num, max_num + 1)
if str(n) not in assigned_nums]
# Also check floating captions that weren't matched
unmatched_floating_nums = {
fc["num"] for fc_idx, fc in enumerate(floating_captions)
if fc_idx not in matched_floating
}
gap_nums = [n for n in gap_nums if str(n) not in unmatched_floating_nums]
for gap_num in gap_nums:
# Determine page range: between pages of adjacent assigned captions
lower_pages = [items[int_nums[n]].page_num
for n in int_nums if n < gap_num]
upper_pages = [items[int_nums[n]].page_num
for n in int_nums if n > gap_num]
page_lo = max(lower_pages) if lower_pages else 1
page_hi = min(upper_pages) if upper_pages else len(doc)
found_caption = _search_page_text_for_caption(
doc, ref_re, gap_num, page_lo, page_hi, kind,
)
if not found_caption:
continue
# Try to match to an orphan on the found page
caption_text, caption_page = found_caption
best_orphan = None
best_dist = float("inf")
for orphan_idx in orphan_indices:
orphan = items[orphan_idx]
if orphan.page_num != caption_page:
continue
dist = orphan_idx # prefer earlier orphan on same page
if dist < best_dist:
best_dist = dist
best_orphan = orphan_idx
if best_orphan is not None:
items[best_orphan].caption = caption_text
assigned_nums[str(gap_num)] = best_orphan
orphan_indices.remove(best_orphan)
recoveries += 1
logger.debug(
"Recovery: %s gap %d filled on p%d",
kind, gap_num, caption_page,
)
return recoveries
def _search_page_text_for_caption(
doc: pymupdf.Document,
ref_re: re.Pattern,
target_num: int,
page_lo: int,
page_hi: int,
kind: str,
) -> tuple[str, int] | None:
"""Search page text for a specific caption number.
Returns (caption_text, page_num) or None.
Rejects body-text references like "as shown in Figure N".
"""
for page_num_0 in range(page_lo - 1, min(page_hi, len(doc))):
page = doc[page_num_0]
page_num = page_num_0 + 1
text_dict = page.get_text("dict")
for block in text_dict.get("blocks", []):
if block.get("type") != 0:
continue
block_text = ""
for line in block.get("lines", []):
for span in line.get("spans", []):
block_text += span.get("text", "")
block_text += " "
block_text = block_text.strip()
if not block_text:
continue
# Only check block start (first 400 chars) for caption pattern
check_window = block_text[:400]
m = ref_re.match(check_window)
if not m:
continue
num = int(m.group(1))
if num != target_num:
continue
# Reject body-text references by checking preceding context
start_pos = m.start()
preceding = block_text[max(0, start_pos - 60):start_pos]
if _BODY_REF_CONTEXT_RE.search(preceding):
continue
return block_text, page_num
return None
