diff --git a/pyproject.toml b/pyproject.toml index 83436a4..b0ad5a7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,7 @@ dev = [ "httpx>=0.28", "anyio>=4.0", "pytest-anyio>=0.0.0", + "scipy>=1.14", ] [build-system] diff --git a/tests/test_alignment.py b/tests/test_alignment.py new file mode 100644 index 0000000..cbc8888 --- /dev/null +++ b/tests/test_alignment.py @@ -0,0 +1,51 @@ +from pathlib import Path + +import numpy as np +from PIL import Image +from scipy.ndimage import binary_dilation + +from unredact.pipeline.rasterize import rasterize_pdf +from unredact.pipeline.ocr import ocr_page +from unredact.pipeline.font_detect import detect_font +from unredact.pipeline.overlay import render_overlay + + +def test_overlay_alignment_score(sample_pdf: Path): + """Measure how well the overlay aligns with the original text. + + Strategy: compare the overlay-only green pixels against the original + dark text pixels. Where they overlap (within a few pixels tolerance), + that's a hit. We want a reasonable hit rate. + """ + pages = rasterize_pdf(sample_pdf, first_page=1, last_page=1) + page = pages[0] + lines = ocr_page(page) + font_match = detect_font(lines, page) + overlay = render_overlay(page, lines, font_match) + + orig_arr = np.array(page.convert("L")) # grayscale + overlay_arr = np.array(overlay) + + # Original dark text pixels (< 128 on grayscale) + text_mask = orig_arr < 128 + + # Overlay green pixels (green channel > 100) + green_mask = overlay_arr[:, :, 1] > 100 + + # How many text pixels have green overlay nearby (within 3px tolerance)? + green_dilated = binary_dilation(green_mask, iterations=3) + + overlap = text_mask & green_dilated + if text_mask.sum() == 0: + return + + hit_rate = overlap.sum() / text_mask.sum() + print(f"\nAlignment hit rate: {hit_rate:.1%}") + print(f" Text pixels: {text_mask.sum()}") + print(f" Green pixels: {green_mask.sum()}") + print(f" Overlap pixels: {overlap.sum()}") + + # We want at least 30% of text pixels to have overlay nearby + # (not 100% because redaction boxes are dark too, and headers/footers + # may use different fonts) + assert hit_rate > 0.3, f"Overlay alignment too low: {hit_rate:.1%}" diff --git a/unredact/pipeline/font_detect.py b/unredact/pipeline/font_detect.py index 5385b2b..45e132f 100644 --- a/unredact/pipeline/font_detect.py +++ b/unredact/pipeline/font_detect.py @@ -100,11 +100,13 @@ def detect_font( ) -> FontMatch: """Detect the best matching font and size for OCR'd text. - Tries all candidate fonts at all sizes in SIZE_RANGE and returns - the combination with the lowest mean word-width error. + Two-pass search: + 1. Coarse pass: all candidate fonts at sizes 20-78 in steps of 2 + 2. Fine pass: best font at sizes ±3 around the coarse best in steps of 1 """ best: FontMatch | None = None + # Coarse pass for font_name in CANDIDATE_FONTS: font_path = _find_font_path(font_name) if font_path is None: @@ -129,4 +131,22 @@ def detect_font( if best is None: raise RuntimeError("No matching font found. Check system fonts.") + # Fine pass: search ±3 around the best size in steps of 1 + fine_range = range(max(8, best.font_size - 3), best.font_size + 4) + for size in fine_range: + if size == best.font_size: + continue + try: + font = ImageFont.truetype(str(best.font_path), size) + except Exception: + continue + score = _score_font(font, lines) + if score < best.score: + best = FontMatch( + font_name=best.font_name, + font_path=best.font_path, + font_size=size, + score=score, + ) + return best diff --git a/unredact/pipeline/overlay.py b/unredact/pipeline/overlay.py index 10371c0..63674be 100644 --- a/unredact/pipeline/overlay.py +++ b/unredact/pipeline/overlay.py @@ -1,7 +1,35 @@ from PIL import Image, ImageDraw, ImageFont from unredact.pipeline.font_detect import FontMatch -from unredact.pipeline.ocr import OcrLine +from unredact.pipeline.ocr import OcrChar, OcrLine + + +def _extract_words(line: OcrLine) -> list[tuple[str, int, int]]: + """Extract words from an OcrLine with their x,y positions. + + Returns list of (word_text, x, y) tuples. + """ + words: list[tuple[str, int, int]] = [] + current_word = "" + word_x = -1 + word_y = -1 + + for char in line.chars: + if char.text == " ": + if current_word: + words.append((current_word, word_x, word_y)) + current_word = "" + word_x = -1 + else: + if not current_word: + word_x = char.x + word_y = char.y + current_word += char.text + + if current_word: + words.append((current_word, word_x, word_y)) + + return words def render_overlay( @@ -12,9 +40,10 @@ def render_overlay( ) -> Image.Image: """Render green text overlay on top of the document image. - Draws each OCR'd line's text at its detected position using - the matched font. The overlay is semi-transparent so the - original document is still visible underneath. + Draws each OCR'd word at its detected position using the matched font. + Words are rendered individually at their OCR'd x positions so that + word spacing matches the original document exactly. Only the internal + character spacing within each word depends on the font metrics. Args: page_image: The original rasterized page. @@ -30,15 +59,17 @@ def render_overlay( draw = ImageDraw.Draw(overlay) font = font_match.to_pil_font() + # Compute y-offset: Pillow draws from the top of the glyph bbox, + # but OCR y is the top of the word bounding box. We need to align these. + glyph_top = font.getbbox("Ag")[1] # typically negative or 0 + for line in lines: if not line.chars: continue - text = line.text - # Position: use the first character's x and the line's y - x = line.x - y = line.y - - draw.text((x, y), text, font=font, fill=color) + words = _extract_words(line) + for word_text, wx, wy in words: + # Adjust y by subtracting the glyph top offset + draw.text((wx, wy - glyph_top), word_text, font=font, fill=color) return Image.alpha_composite(base, overlay)