diff --git a/tests/test_analyze_page.py b/tests/test_analyze_page.py new file mode 100644 index 0000000..afcf4ac --- /dev/null +++ b/tests/test_analyze_page.py @@ -0,0 +1,317 @@ +"""Tests for the page analysis pipeline (analyze_page).""" + +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import numpy as np +import pytest +from PIL import Image + +from unredact.pipeline.ocr import OcrChar, OcrLine +from unredact.pipeline.llm_detect import LlmRedaction +from unredact.pipeline.detect_redactions import Redaction +from unredact.pipeline.font_detect import FontMatch +from unredact.pipeline.analyze_page import ( + RedactionAnalysis, + PageAnalysis, + analyze_page, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_char(ch: str, x: int, w: int = 10, y: int = 100, h: int = 20) -> OcrChar: + return OcrChar(text=ch, x=x, y=y, w=w, h=h, conf=95.0) + + +def _make_line_from_words( + words: list[str], + start_x: int = 50, + char_w: int = 10, + space_w: int = 8, + y: int = 100, + h: int = 20, +) -> OcrLine: + """Build an OcrLine from a list of words with evenly-spaced characters.""" + chars: list[OcrChar] = [] + x = start_x + for wi, word in enumerate(words): + for ch in word: + chars.append(_make_char(ch, x, w=char_w, y=y, h=h)) + x += char_w + if wi < len(words) - 1: + chars.append(_make_char(" ", x, w=space_w, y=y, h=h)) + x += space_w + line_x = chars[0].x + line_w = (chars[-1].x + chars[-1].w) - line_x + return OcrLine(chars=chars, x=line_x, y=y, w=line_w, h=h) + + +def _dummy_page_image(width: int = 800, height: int = 600) -> Image.Image: + """Create a white test image.""" + return Image.new("RGB", (width, height), "white") + + +def _dummy_font_match() -> FontMatch: + from unredact.pipeline.font_detect import _find_font_path + + font_path = _find_font_path("Times New Roman") + if font_path is None: + font_path = _find_font_path("serif") + if font_path is None: + font_path = Path("/usr/share/fonts/TTF/Times.TTF") + + return FontMatch( + font_name="Times New Roman", + font_path=font_path, + font_size=20, + score=0.85, + ) + + +# --------------------------------------------------------------------------- +# test_analyze_page_returns_results +# --------------------------------------------------------------------------- + +@pytest.mark.anyio +async def test_analyze_page_returns_results(): + """Mock LLM returning 1 redaction, OpenCV returning a box, font detection + returning a FontMatch. Verify PageAnalysis has 1 RedactionAnalysis with + correct fields.""" + line = _make_line_from_words(["Hello", "|||", "world"], y=100, h=20) + lines = [line] + + llm_redaction = LlmRedaction( + line_index=0, + left_word="Hello", + right_word="world", + left_x=100, + right_x=130, + line_y=100, + line_h=20, + ) + + redaction_box = Redaction( + id="abc12345", + x=100, + y=98, + w=30, + h=22, + ) + + font_match = _dummy_font_match() + page_image = _dummy_page_image() + + with ( + patch( + "unredact.pipeline.analyze_page.ocr_page", + return_value=lines, + ), + patch( + "unredact.pipeline.analyze_page.detect_redactions_llm", + new_callable=AsyncMock, + return_value=[llm_redaction], + ), + patch( + "unredact.pipeline.analyze_page.find_redaction_in_region", + return_value=redaction_box, + ), + patch( + "unredact.pipeline.analyze_page.detect_font_masked", + return_value=font_match, + ), + patch( + "unredact.pipeline.analyze_page.align_text_to_page", + return_value=(3, -1), + ), + ): + result = await analyze_page(page_image) + + assert isinstance(result, PageAnalysis) + assert len(result.lines) == 1 + assert len(result.redactions) == 1 + + ra = result.redactions[0] + assert isinstance(ra, RedactionAnalysis) + assert ra.box is redaction_box + assert ra.line is line + assert ra.font is font_match + + # left_text: chars whose center is left of box.x (100) + # "Hello" chars: x=50..90 (centers at 55, 65, 75, 85, 95) — all < 100 + # Space char: x=100, center=104 — not < 100 + assert "Hello" in ra.left_text + + # right_text: chars whose center is right of box.x+box.w (130) + # "world" starts after the space and ||| + assert "world" in ra.right_text + + +# --------------------------------------------------------------------------- +# test_analyze_page_no_redactions +# --------------------------------------------------------------------------- + +@pytest.mark.anyio +async def test_analyze_page_no_redactions(): + """Mock LLM returning empty list. Verify PageAnalysis has empty redactions.""" + line = _make_line_from_words(["Clean", "text", "here"]) + lines = [line] + page_image = _dummy_page_image() + + with ( + patch( + "unredact.pipeline.analyze_page.ocr_page", + return_value=lines, + ), + patch( + "unredact.pipeline.analyze_page.detect_redactions_llm", + new_callable=AsyncMock, + return_value=[], + ), + ): + result = await analyze_page(page_image) + + assert isinstance(result, PageAnalysis) + assert len(result.lines) == 1 + assert result.redactions == [] + + +# --------------------------------------------------------------------------- +# test_analyze_page_skips_unfound_boxes +# --------------------------------------------------------------------------- + +@pytest.mark.anyio +async def test_analyze_page_skips_unfound_boxes(): + """Mock LLM returning 1 redaction but OpenCV returns None. + Verify it's skipped gracefully.""" + line = _make_line_from_words(["The", "|||", "cat"]) + lines = [line] + + llm_redaction = LlmRedaction( + line_index=0, + left_word="The", + right_word="cat", + left_x=80, + right_x=110, + line_y=100, + line_h=20, + ) + + page_image = _dummy_page_image() + + with ( + patch( + "unredact.pipeline.analyze_page.ocr_page", + return_value=lines, + ), + patch( + "unredact.pipeline.analyze_page.detect_redactions_llm", + new_callable=AsyncMock, + return_value=[llm_redaction], + ), + patch( + "unredact.pipeline.analyze_page.find_redaction_in_region", + return_value=None, + ), + ): + result = await analyze_page(page_image) + + assert isinstance(result, PageAnalysis) + assert len(result.lines) == 1 + assert result.redactions == [] + + +# --------------------------------------------------------------------------- +# test_analyze_page_progress_callback +# --------------------------------------------------------------------------- + +@pytest.mark.anyio +async def test_analyze_page_progress_callback(): + """Verify that on_progress is called at each stage.""" + line = _make_line_from_words(["Hello", "world"]) + lines = [line] + page_image = _dummy_page_image() + + progress_events: list[tuple] = [] + + def on_progress(event: str, data: dict): + progress_events.append((event, data)) + + with ( + patch( + "unredact.pipeline.analyze_page.ocr_page", + return_value=lines, + ), + patch( + "unredact.pipeline.analyze_page.detect_redactions_llm", + new_callable=AsyncMock, + return_value=[], + ), + ): + await analyze_page(page_image, on_progress=on_progress) + + event_names = [e[0] for e in progress_events] + assert "ocr_done" in event_names + assert "redactions_found" in event_names + assert "analysis_complete" in event_names + + +# --------------------------------------------------------------------------- +# test_analyze_page_font_caching +# --------------------------------------------------------------------------- + +@pytest.mark.anyio +async def test_analyze_page_font_caching(): + """Two redactions on the same line should call detect_font_masked only once + for that line (font caching).""" + line = _make_line_from_words( + ["A", "XXX", "B", "YYY", "C"], start_x=50, y=100, h=20, + ) + lines = [line] + + llm_red_1 = LlmRedaction( + line_index=0, left_word="A", right_word="B", + left_x=60, right_x=80, line_y=100, line_h=20, + ) + llm_red_2 = LlmRedaction( + line_index=0, left_word="B", right_word="C", + left_x=90, right_x=110, line_y=100, line_h=20, + ) + + box1 = Redaction(id="b1", x=60, y=98, w=20, h=22) + box2 = Redaction(id="b2", x=90, y=98, w=20, h=22) + + font_match = _dummy_font_match() + page_image = _dummy_page_image() + + mock_font_detect = patch( + "unredact.pipeline.analyze_page.detect_font_masked", + return_value=font_match, + ) + with ( + patch("unredact.pipeline.analyze_page.ocr_page", return_value=lines), + patch( + "unredact.pipeline.analyze_page.detect_redactions_llm", + new_callable=AsyncMock, + return_value=[llm_red_1, llm_red_2], + ), + patch( + "unredact.pipeline.analyze_page.find_redaction_in_region", + side_effect=[box1, box2], + ), + mock_font_detect as mock_fd, + patch( + "unredact.pipeline.analyze_page.align_text_to_page", + return_value=(0, 0), + ), + ): + result = await analyze_page(page_image) + + assert len(result.redactions) == 2 + # Font detection should have been called only once for line 0 + assert mock_fd.call_count == 1 + # But both redactions should share the same font + assert result.redactions[0].font is font_match + assert result.redactions[1].font is font_match diff --git a/unredact/pipeline/analyze_page.py b/unredact/pipeline/analyze_page.py new file mode 100644 index 0000000..cf397ee --- /dev/null +++ b/unredact/pipeline/analyze_page.py @@ -0,0 +1,192 @@ +"""Page analysis pipeline: OCR -> LLM -> guided OpenCV -> font detection. + +Orchestrates the full per-page pipeline, wiring together OCR, LLM-based +redaction detection, guided OpenCV search, and font detection with masking. +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass + +import numpy as np +from PIL import Image + +from unredact.pipeline.ocr import ocr_page, OcrLine +from unredact.pipeline.llm_detect import detect_redactions_llm, LlmRedaction +from unredact.pipeline.detect_redactions import find_redaction_in_region, Redaction +from unredact.pipeline.font_detect import detect_font_masked, align_text_to_page, FontMatch + +log = logging.getLogger(__name__) + + +@dataclass +class RedactionAnalysis: + """Analysis result for a single redaction on a page.""" + + box: Redaction # Pixel-precise bounding box + line: OcrLine # The OCR line containing this redaction + font: FontMatch # Detected font for this line + left_text: str # Clean text to the left of redaction + right_text: str # Clean text to the right of redaction + offset_x: float # Pixel alignment offset X + offset_y: float # Pixel alignment offset Y + + +@dataclass +class PageAnalysis: + """Full analysis result for a single page.""" + + lines: list[OcrLine] # All OCR lines on the page + redactions: list[RedactionAnalysis] # Analysis for each redaction + + +async def analyze_page( + page_image: Image.Image, + on_progress: callable | None = None, +) -> PageAnalysis: + """Run the full per-page analysis pipeline. + + Pipeline steps: + 1. OCR the page image (blocking, run in thread) + 2. LLM detection of redactions (async) + 3. For each LLM redaction: + a. Guided OpenCV search for the precise box + b. Font detection with redaction masking (cached per line) + c. Extract left/right text segments around the redaction + d. Pixel alignment of rendered text to page + + Args: + page_image: PIL Image of the document page. + on_progress: Optional callback ``(event_name, data_dict)`` for + progress reporting. + + Returns: + PageAnalysis with OCR lines and redaction analysis results. + """ + # Step 1: OCR (blocking — run in thread) + lines: list[OcrLine] = await asyncio.to_thread(ocr_page, page_image) + + if on_progress: + on_progress("ocr_done", {"line_count": len(lines)}) + + # Step 2: LLM detection (async) + llm_redactions: list[LlmRedaction] = await detect_redactions_llm(lines) + + if on_progress: + on_progress("redactions_found", {"count": len(llm_redactions)}) + + # Step 3: Process each LLM redaction + # Group LLM redactions by line index so we can collect all boxes per line + # before calling font detection (which needs all boxes for masking). + line_llm_reds: dict[int, list[LlmRedaction]] = {} + for llm_red in llm_redactions: + line_llm_reds.setdefault(llm_red.line_index, []).append(llm_red) + + # 3a: Guided OpenCV for each LLM redaction — build mapping from + # LlmRedaction to Redaction box (or None if not found). + llm_to_box: dict[int, Redaction | None] = {} + for i, llm_red in enumerate(llm_redactions): + box = await asyncio.to_thread( + find_redaction_in_region, + page_image, + llm_red.left_x, + llm_red.line_y, + llm_red.right_x, + llm_red.line_y + llm_red.line_h, + ) + llm_to_box[i] = box + + # 3b: Font detection with caching per line index. + # Pass ALL redaction boxes on the line for masking. + line_fonts: dict[int, FontMatch] = {} + + for line_idx, reds_on_line in line_llm_reds.items(): + # Collect all successfully found boxes on this line + redaction_boxes: list[tuple[int, int, int, int]] = [] + for llm_red in reds_on_line: + idx_in_all = llm_redactions.index(llm_red) + box = llm_to_box[idx_in_all] + if box is not None: + redaction_boxes.append((box.x, box.y, box.w, box.h)) + + # Only detect font if we have at least one valid box on this line + if not redaction_boxes: + continue + + if 0 <= line_idx < len(lines): + line = lines[line_idx] + font = await asyncio.to_thread( + detect_font_masked, line, page_image, redaction_boxes, + ) + line_fonts[line_idx] = font + + # 3c-3e: Build RedactionAnalysis objects + results: list[RedactionAnalysis] = [] + for i, llm_red in enumerate(llm_redactions): + box = llm_to_box[i] + if box is None: + log.info( + "OpenCV found no box for LLM redaction on line %d, skipping", + llm_red.line_index, + ) + continue + + if llm_red.line_index < 0 or llm_red.line_index >= len(lines): + log.warning("Invalid line_index %d, skipping", llm_red.line_index) + continue + + line = lines[llm_red.line_index] + font = line_fonts.get(llm_red.line_index) + if font is None: + log.warning( + "No font detected for line %d, skipping", llm_red.line_index, + ) + continue + + # 3d: Extract left/right text using center-point char filtering + left_chars = [ + c for c in line.chars if c.x + c.w / 2 < box.x + ] + right_chars = [ + c for c in line.chars if c.x + c.w / 2 > box.x + box.w + ] + left_text = "".join(c.text for c in left_chars).strip() + right_text = "".join(c.text for c in right_chars).strip() + + # 3e: Pixel alignment + offset_x = 0.0 + offset_y = 0.0 + + if left_text: + pil_font = font.to_pil_font() + text_region_x1 = max(0, line.x - 20) + text_region_x2 = min(page_image.width, box.x + 20) + text_region_y1 = max(0, line.y - 10) + text_region_y2 = min(page_image.height, line.y + line.h + 10) + text_crop = np.array(page_image.convert("L").crop( + (text_region_x1, text_region_y1, text_region_x2, text_region_y2) + )) + align_dx, align_dy = align_text_to_page( + left_text, pil_font, text_crop, + ) + offset_x = float(text_region_x1 + align_dx - line.x) + offset_y = float(text_region_y1 + align_dy - line.y) + + results.append( + RedactionAnalysis( + box=box, + line=line, + font=font, + left_text=left_text, + right_text=right_text, + offset_x=offset_x, + offset_y=offset_y, + ) + ) + + if on_progress: + on_progress("analysis_complete", {"count": len(results)}) + + return PageAnalysis(lines=lines, redactions=results)