diff --git a/tests/test_font_masking.py b/tests/test_font_masking.py new file mode 100644 index 0000000..80e0b48 --- /dev/null +++ b/tests/test_font_masking.py @@ -0,0 +1,155 @@ +"""Test that detect_font_masked correctly identifies fonts even with redaction boxes.""" +import numpy as np +from PIL import Image, ImageDraw, ImageFont + +from unredact.pipeline.font_detect import ( + _find_font_path, + detect_font_masked, +) +from unredact.pipeline.ocr import OcrChar, OcrLine + + +def _render_page_with_redaction( + font_name: str, + font_size: int, + text: str, + redaction_frac: tuple[float, float] = (0.4, 0.7), +) -> tuple[OcrLine, Image.Image, list[tuple[int, int, int, int]]]: + """Render text onto a page image with a black redaction box. + + Returns (OcrLine, page_image, redaction_boxes) where redaction_boxes + are (rx, ry, rw, rh) in page-relative coordinates. + + The OcrLine bounding box is set to tightly wrap the rendered text + (mimicking what OCR would produce), so detect_font_masked can create + a scoring_line with x=0, y=0 that aligns correctly. + """ + font_path = _find_font_path(font_name) + assert font_path is not None, f"Font {font_name} not found on system" + font = ImageFont.truetype(str(font_path), font_size) + + # Measure the text + bbox = font.getbbox(text) + text_w = bbox[2] - bbox[0] + text_h = bbox[3] - bbox[1] + + # Place the line at an offset on a larger "page" image. + # The line bounding box starts exactly where the text ink begins, + # matching what OCR would detect. + page_margin_x = 50 + page_margin_y = 80 + + # Text is rendered at (render_x, render_y) such that ink starts at + # (line_x, line_y) in page coordinates. + line_x = page_margin_x + line_y = page_margin_y + render_x = line_x - bbox[0] + render_y = line_y - bbox[1] + line_w = text_w + line_h = text_h + + page_w = line_x + line_w + 50 + page_h = line_y + line_h + 100 + + # Render the text onto the page + page_img = Image.new("L", (page_w, page_h), 255) + draw = ImageDraw.Draw(page_img) + draw.text((render_x, render_y), text, font=font, fill=0) + + # Build an OcrLine with approximate character positions (page-relative) + words = text.split() + chars: list[OcrChar] = [] + cursor_x = line_x + for wi, word in enumerate(words): + word_bbox = font.getbbox(word) + word_w = word_bbox[2] - word_bbox[0] + char_w = word_w / len(word) + for ci, ch in enumerate(word): + chars.append(OcrChar( + text=ch, + x=int(cursor_x + ci * char_w), + y=line_y, + w=max(1, int(char_w)), + h=text_h, + conf=95.0, + )) + cursor_x += word_w + if wi < len(words) - 1: + space_w = font.getlength(" ") + chars.append(OcrChar( + text=" ", + x=int(cursor_x), + y=line_y, + w=max(1, int(space_w)), + h=text_h, + conf=95.0, + )) + cursor_x += space_w + + line = OcrLine( + chars=chars, + x=line_x, + y=line_y, + w=line_w, + h=line_h, + ) + + # Draw a black redaction box covering a fraction of the line + redact_x = line_x + int(line_w * redaction_frac[0]) + redact_y = line_y + redact_w = int(line_w * (redaction_frac[1] - redaction_frac[0])) + redact_h = line_h + draw.rectangle( + [redact_x, redact_y, redact_x + redact_w, redact_y + redact_h], + fill=0, + ) + + redaction_boxes = [(redact_x, redact_y, redact_w, redact_h)] + + # Convert to RGB page image (detect_font_masked expects PIL Image) + page_rgb = page_img.convert("RGB") + return line, page_rgb, redaction_boxes + + +def test_masked_detection_finds_correct_font(): + """Render TNR@50, add a black redaction box, and verify detection still finds TNR.""" + text = "On Aug 27 2012 at 12:52 PM wrote ot" + line, page_img, redaction_boxes = _render_page_with_redaction( + font_name="Times New Roman", + font_size=50, + text=text, + ) + + result = detect_font_masked(line, page_img, redaction_boxes) + + assert "Times" in result.font_name or "Liberation Serif" in result.font_name, ( + f"Expected serif font, got {result.font_name}" + ) + assert 47 <= result.font_size <= 53, ( + f"Expected size ~50, got {result.font_size}" + ) + assert result.score > 0.3, ( + f"Expected reasonable score, got {result.score:.3f}" + ) + + +def test_masked_detection_not_confused_by_redaction(): + """Render Arial@40, add a black redaction box, and verify it finds Arial.""" + text = "On Aug 27 2012 at 12:52 PM wrote ot" + line, page_img, redaction_boxes = _render_page_with_redaction( + font_name="Arial", + font_size=40, + text=text, + ) + + result = detect_font_masked(line, page_img, redaction_boxes) + + assert "Arial" in result.font_name or "Liberation Sans" in result.font_name, ( + f"Expected sans-serif font, got {result.font_name}" + ) + assert 37 <= result.font_size <= 43, ( + f"Expected size ~40, got {result.font_size}" + ) + assert result.score > 0.3, ( + f"Expected reasonable score, got {result.score:.3f}" + ) diff --git a/unredact/pipeline/font_detect.py b/unredact/pipeline/font_detect.py index ac1a34f..7e96a80 100644 --- a/unredact/pipeline/font_detect.py +++ b/unredact/pipeline/font_detect.py @@ -274,6 +274,62 @@ def detect_font_for_line_from_crop( return _fine_search(line, line_crop, best) +def detect_font_masked( + line: OcrLine, + page_image: Image.Image, + redaction_boxes: list[tuple[int, int, int, int]], +) -> FontMatch: + """Detect the best font by masking redaction boxes to white before scoring. + + Instead of cropping just the unredacted portion of a line (which loses + signal from narrow crops and coordinate mismatches), this masks each + redaction box to white (255) in the full line crop. The Dice-based pixel + scoring naturally ignores the masked (white) regions since they contain + no ink pixels. + + Args: + line: OCR line with page-relative coordinates. + page_image: Full page as a PIL Image (any mode; converted to grayscale). + redaction_boxes: List of (rx, ry, rw, rh) rectangles in page-relative + coordinates. Each rectangle is painted white before scoring. + + Returns: + FontMatch with the best font name, path, size, and score. + """ + # 1. Crop the line region from the page image + line_crop = np.array( + page_image.convert("L").crop((line.x, line.y, line.x + line.w, line.y + line.h)) + ) + + # 2. Mask each redaction box to white (255) — convert page-relative to crop-relative + for rx, ry, rw, rh in redaction_boxes: + # Convert from page coordinates to crop-relative coordinates + cx = rx - line.x + cy = ry - line.y + + # Clip to the crop bounds + x0 = max(0, cx) + y0 = max(0, cy) + x1 = min(line.w, cx + rw) + y1 = min(line.h, cy + rh) + + if x0 < x1 and y0 < y1: + line_crop[y0:y1, x0:x1] = 255 + + # 3. Create a scoring line with crop-relative coordinates (x=0, y=0) + scoring_line = OcrLine( + chars=line.chars, + x=0, y=0, + w=line.w, h=line.h, + ) + + # 4. Full search then fine search + best = _full_search(scoring_line, line_crop) + if best is None: + raise RuntimeError("No matching font found. Check system fonts.") + return _fine_search(scoring_line, line_crop, best) + + def detect_font_for_line( line: OcrLine, page_image: Image.Image,