Spaces:
Sleeping
Sleeping
File size: 2,949 Bytes
0e39d80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | """PII masking using OCR bounding boxes — no YOLO dependency."""
from pathlib import Path
import re
import cv2
import numpy as np
from PIL import Image, ImageDraw
from ml_utils.ocr import OcrResult
# Regex patterns for PII fields to mask
PII_PATTERNS: dict[str, list[re.Pattern]] = {
"aadhaar": [
re.compile(r"\d{4}\s?\d{4}\s?\d{4}"), # Full Aadhaar number
],
"pan": [
re.compile(r"[A-Z]{5}\d{4}[A-Z]"), # PAN number
],
}
# Additional keyword-based masking: if an OCR block near these labels,
# mask the VALUE block (the one to the right/below)
PII_LABELS: dict[str, list[str]] = {
"aadhaar": ["address", "पता"],
"pan": [],
}
def mask_pii_on_image(
image_bgr: np.ndarray,
ocr_results: list[OcrResult],
doc_type: str,
) -> np.ndarray:
"""Mask PII fields on the image using OCR bounding boxes."""
img_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
img_pil = Image.fromarray(img_rgb)
draw = ImageDraw.Draw(img_pil)
patterns = PII_PATTERNS.get(doc_type, [])
for result in ocr_results:
for pat in patterns:
if pat.search(result.text):
_draw_mask(draw, result.bbox, img_pil.size)
break
# Mask address region for Aadhaar (bottom-right blocks after "Address" label)
if doc_type == "aadhaar":
_mask_address_region(draw, ocr_results, img_pil.size)
return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
def _draw_mask(draw: ImageDraw.ImageDraw, bbox: list[list[int]], img_size: tuple[int, int]):
"""Draw a black rectangle over a bounding box."""
w, h = img_size
xs = [max(0, min(p[0], w)) for p in bbox]
ys = [max(0, min(p[1], h)) for p in bbox]
x1, x2 = min(xs), max(xs)
y1, y2 = min(ys), max(ys)
# Add small padding
pad = 3
draw.rectangle([x1 - pad, y1 - pad, x2 + pad, y2 + pad], fill=(0, 0, 0))
def _mask_address_region(draw: ImageDraw.ImageDraw, results: list[OcrResult], img_size: tuple[int, int]):
"""Mask text blocks in the address region of an Aadhaar card."""
# Find "Address" or "पता" label
address_label = None
for r in results:
if any(kw in r.text.lower() for kw in ["address", "पता"]):
address_label = r
break
if address_label is None:
return
# Mask all blocks below the address label (within reasonable distance)
_, label_y1, _, label_y2 = address_label.rect
img_h = img_size[1]
max_y = label_y2 + (img_h - label_y2) # everything below
for r in results:
_, ry1, _, _ = r.rect
if ry1 >= label_y1 and r is not address_label:
_draw_mask(draw, r.bbox, img_size)
def save_masked_image(image_bgr: np.ndarray, output_path: Path) -> str:
"""Save masked image to disk."""
output_path.parent.mkdir(parents=True, exist_ok=True)
cv2.imwrite(str(output_path), image_bgr)
return str(output_path)
|