diff --git a/pyproject.toml b/pyproject.toml index 299ea51..8fbf919 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "opencv-python>=4.10", "Pillow>=11.0", "anthropic>=0.52", + "numpy>=1.26", ] [project.optional-dependencies] diff --git a/tests/test_overlay.py b/tests/test_overlay.py new file mode 100644 index 0000000..f9e8bd6 --- /dev/null +++ b/tests/test_overlay.py @@ -0,0 +1,41 @@ +from pathlib import Path + +import numpy as np +from PIL import Image + +from unredact.pipeline.overlay import render_overlay +from unredact.pipeline.ocr import ocr_page +from unredact.pipeline.font_detect import detect_font +from unredact.pipeline.rasterize import rasterize_pdf + + +def test_overlay_same_size_as_original(sample_pdf: Path): + pages = rasterize_pdf(sample_pdf, first_page=1, last_page=1) + lines = ocr_page(pages[0]) + font_match = detect_font(lines, pages[0]) + result = render_overlay(pages[0], lines, font_match) + assert result.size == pages[0].size + + +def test_overlay_has_green_pixels(sample_pdf: Path): + pages = rasterize_pdf(sample_pdf, first_page=1, last_page=1) + lines = ocr_page(pages[0]) + font_match = detect_font(lines, pages[0]) + result = render_overlay(pages[0], lines, font_match) + # Convert to numpy and check for green-ish pixels + arr = np.array(result) + # Green channel significantly higher than red and blue + green_mask = (arr[:, :, 1] > 100) & (arr[:, :, 0] < 100) & (arr[:, :, 2] < 100) + green_pixel_count = green_mask.sum() + assert green_pixel_count > 100, "Expected green text pixels in the overlay" + + +def test_overlay_preserves_original(sample_pdf: Path): + """Overlay should be composited — original should still be visible.""" + pages = rasterize_pdf(sample_pdf, first_page=1, last_page=1) + original = pages[0].copy() + lines = ocr_page(pages[0]) + font_match = detect_font(lines, pages[0]) + result = render_overlay(pages[0], lines, font_match) + # The result shouldn't be identical to the original (overlay adds green) + assert list(result.getdata()) != list(original.getdata()) diff --git a/unredact/pipeline/overlay.py b/unredact/pipeline/overlay.py new file mode 100644 index 0000000..10371c0 --- /dev/null +++ b/unredact/pipeline/overlay.py @@ -0,0 +1,44 @@ +from PIL import Image, ImageDraw, ImageFont + +from unredact.pipeline.font_detect import FontMatch +from unredact.pipeline.ocr import OcrLine + + +def render_overlay( + page_image: Image.Image, + lines: list[OcrLine], + font_match: FontMatch, + color: tuple[int, int, int, int] = (0, 200, 0, 160), +) -> 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. + + Args: + page_image: The original rasterized page. + lines: OCR'd lines with character positions. + font_match: Detected font to use for rendering. + color: RGBA color for the overlay text. + + Returns: + A new RGBA image with the overlay composited on the original. + """ + base = page_image.convert("RGBA") + overlay = Image.new("RGBA", base.size, (0, 0, 0, 0)) + draw = ImageDraw.Draw(overlay) + font = font_match.to_pil_font() + + 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) + + return Image.alpha_composite(base, overlay)