diff --git a/tests/test_detect_redactions.py b/tests/test_detect_redactions.py new file mode 100644 index 0000000..d4f6101 --- /dev/null +++ b/tests/test_detect_redactions.py @@ -0,0 +1,52 @@ +import numpy as np +from PIL import Image + +from unredact.pipeline.detect_redactions import detect_redactions, Redaction + + +def test_detect_single_black_rectangle(): + """A white image with one black rectangle should yield one redaction.""" + img = Image.new("RGB", (800, 600), "white") + pixels = np.array(img) + # Draw a black rectangle (simulating a redaction) + pixels[200:230, 100:300] = [0, 0, 0] + img = Image.fromarray(pixels) + + redactions = detect_redactions(img) + assert len(redactions) == 1 + r = redactions[0] + assert abs(r.x - 100) < 5 + assert abs(r.y - 200) < 5 + assert abs(r.w - 200) < 10 + assert abs(r.h - 30) < 10 + + +def test_detect_no_redactions_on_clean_page(): + """A white page with no black rectangles should return empty list.""" + img = Image.new("RGB", (800, 600), "white") + redactions = detect_redactions(img) + assert redactions == [] + + +def test_detect_ignores_small_noise(): + """Small black spots (< min area) should be ignored.""" + img = Image.new("RGB", (800, 600), "white") + pixels = np.array(img) + pixels[100:105, 100:105] = [0, 0, 0] # 5x5 spot — too small + img = Image.fromarray(pixels) + + redactions = detect_redactions(img) + assert redactions == [] + + +def test_detect_multiple_redactions(): + """Multiple black rectangles should all be detected.""" + img = Image.new("RGB", (800, 600), "white") + pixels = np.array(img) + pixels[100:125, 50:250] = [0, 0, 0] + pixels[300:325, 100:400] = [0, 0, 0] + pixels[450:475, 200:500] = [0, 0, 0] + img = Image.fromarray(pixels) + + redactions = detect_redactions(img) + assert len(redactions) == 3 diff --git a/unredact/pipeline/detect_redactions.py b/unredact/pipeline/detect_redactions.py new file mode 100644 index 0000000..fddebd8 --- /dev/null +++ b/unredact/pipeline/detect_redactions.py @@ -0,0 +1,74 @@ +from dataclasses import dataclass +import uuid + +import cv2 +import numpy as np +from PIL import Image + + +@dataclass +class Redaction: + """A detected redaction bounding box.""" + id: str + x: int + y: int + w: int + h: int + + +# Minimum area in pixels to consider (filters noise) +MIN_AREA = 500 + +# Minimum aspect ratio (width/height) — redactions are wider than tall +MIN_ASPECT = 1.5 + + +def detect_redactions(image: Image.Image) -> list[Redaction]: + """Detect black-filled rectangles in a page image. + + Converts to grayscale, thresholds for near-black pixels, finds contours, + and filters for rectangular shapes that look like redaction bars. + + Args: + image: PIL Image of a document page. + + Returns: + List of Redaction objects sorted top-to-bottom, left-to-right. + """ + arr = np.array(image.convert("RGB")) + gray = cv2.cvtColor(arr, cv2.COLOR_RGB2GRAY) + + # Threshold: pixels darker than 40 are "black" + _, binary = cv2.threshold(gray, 40, 255, cv2.THRESH_BINARY_INV) + + # Morphological close to merge adjacent redaction fragments + kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 3)) + binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel) + + contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + redactions: list[Redaction] = [] + for contour in contours: + x, y, w, h = cv2.boundingRect(contour) + area = w * h + + if area < MIN_AREA: + continue + + # Aspect ratio filter: redactions are wider than tall + if h > 0 and w / h < MIN_ASPECT: + continue + + # Fill ratio: the contour should fill most of the bounding rect + contour_area = cv2.contourArea(contour) + if contour_area / area < 0.7: + continue + + redactions.append(Redaction( + id=uuid.uuid4().hex[:8], + x=x, y=y, w=w, h=h, + )) + + # Sort top-to-bottom, then left-to-right + redactions.sort(key=lambda r: (r.y, r.x)) + return redactions