"""Face detection and blurring using MediaPipe.""" import sys import json import os def emit_progress(percent, stage): """Emit structured progress to stderr for bridge.ts to capture.""" print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True) # ── Model path for new mp.tasks API ───────────────────────────────── _FACE_DETECT_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.tflite" _DOCKER_MODEL_PATH = "/opt/models/mediapipe/blaze_face_short_range.tflite" _LOCAL_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models") _LOCAL_MODEL_PATH = os.path.join(_LOCAL_MODEL_DIR, "blaze_face_short_range.tflite") def _ensure_face_detect_model(): """Resolve face detector model. Docker path first, then local dev.""" if os.path.exists(_DOCKER_MODEL_PATH): return _DOCKER_MODEL_PATH if os.path.exists(_LOCAL_MODEL_PATH): return _LOCAL_MODEL_PATH os.makedirs(_LOCAL_MODEL_DIR, exist_ok=True) import urllib.request emit_progress(15, "Downloading face detection model") urllib.request.urlretrieve(_FACE_DETECT_MODEL_URL, _LOCAL_MODEL_PATH) return _LOCAL_MODEL_PATH def _iou(a, b): """Compute intersection-over-union between two face boxes.""" ax2, ay2 = a["x"] + a["w"], a["y"] + a["h"] bx2, by2 = b["x"] + b["w"], b["y"] + b["h"] inter_w = max(0, min(ax2, bx2) - max(a["x"], b["x"])) inter_h = max(0, min(ay2, by2) - max(a["y"], b["y"])) inter = inter_w * inter_h union = a["w"] * a["h"] + b["w"] * b["h"] - inter return inter / union if union > 0 else 0.0 def _nms_faces(faces, iou_threshold=0.4): """Remove duplicate detections using greedy non-maximum suppression.""" if len(faces) <= 1: return faces kept = [] used = [False] * len(faces) for i in range(len(faces)): if used[i]: continue kept.append(faces[i]) used[i] = True for j in range(i + 1, len(faces)): if not used[j] and _iou(faces[i], faces[j]) >= iou_threshold: used[j] = True return kept def _detect_with_solutions(img_array, min_confidence): """Detect faces using legacy mp.solutions API (mediapipe < 0.10.30). Runs both short-range (model 0) and full-range (model 1) detectors and merges the results. Previously the loop broke on the first model that found any face, so group photos where model 0 caught only 1-2 large faces would never have the remaining faces scanned by model 1. """ import mediapipe as mp mp_face = mp.solutions.face_detection ih, iw = img_array.shape[:2] all_faces = [] for model_sel in [0, 1]: detector = mp_face.FaceDetection( model_selection=model_sel, min_detection_confidence=min_confidence, ) results = detector.process(img_array) detector.close() for detection in (results.detections or []): bbox = detection.location_data.relative_bounding_box all_faces.append({ "x": int(bbox.xmin * iw), "y": int(bbox.ymin * ih), "w": int(bbox.width * iw), "h": int(bbox.height * ih), }) return _nms_faces(all_faces) def _detect_with_tasks(img_array, min_confidence): """Detect faces using new mp.tasks API (mediapipe >= 0.10.30).""" import mediapipe as mp model_path = _ensure_face_detect_model() options = mp.tasks.vision.FaceDetectorOptions( base_options=mp.tasks.BaseOptions(model_asset_path=model_path), running_mode=mp.tasks.vision.RunningMode.IMAGE, min_detection_confidence=min_confidence, ) detector = mp.tasks.vision.FaceDetector.create_from_options(options) mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=img_array) result = detector.detect(mp_image) detector.close() faces = [] for detection in result.detections: bbox = detection.bounding_box faces.append({ "x": bbox.origin_x, "y": bbox.origin_y, "w": bbox.width, "h": bbox.height, }) return faces def _detect_faces(img_array, min_confidence): """Detect faces, trying legacy API first then falling back to tasks API.""" try: return _detect_with_solutions(img_array, min_confidence) except AttributeError: return _detect_with_tasks(img_array, min_confidence) def main(): input_path = sys.argv[1] output_path = sys.argv[2] settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {} blur_radius = settings.get("blurRadius", 30) sensitivity = settings.get("sensitivity", 0.5) detect_only = settings.get("detectOnly", False) try: emit_progress(10, "Preparing") from PIL import Image, ImageFilter img = Image.open(input_path).convert("RGB") try: import numpy as np emit_progress(20, "Ready") # Map sensitivity (0.1-0.9) to MediaPipe confidence threshold. # Higher sensitivity = lower confidence threshold = more detections. min_confidence = max(0.1, 1.0 - sensitivity) img_array = np.array(img) # Try legacy mp.solutions API first, fall back to mp.tasks emit_progress(25, "Scanning for faces") faces = _detect_faces(img_array, min_confidence) num_faces = len(faces) emit_progress(50, f"Found {num_faces} face{'s' if num_faces != 1 else ''}") if num_faces > 0 and not detect_only: for i, face in enumerate(faces): x, y, w, h = face["x"], face["y"], face["w"], face["h"] # Add padding around the face pad = int(max(w, h) * 0.1) x1 = max(0, x - pad) y1 = max(0, y - pad) x2 = min(img.width, x + w + pad) y2 = min(img.height, y + h + pad) face_region = img.crop((x1, y1, x2, y2)) blurred = face_region.filter( ImageFilter.GaussianBlur(blur_radius) ) img.paste(blurred, (x1, y1)) emit_progress( 50 + int((i + 1) / num_faces * 40), f"Blurring face {i + 1} of {num_faces}", ) if not detect_only: emit_progress(95, "Saving result") img.save(output_path) print( json.dumps( { "success": True, "facesDetected": len(faces), "faces": faces, } ) ) except ImportError: print( json.dumps( { "success": False, "error": "Face detection requires MediaPipe. Install with: pip install mediapipe", } ) ) sys.exit(1) except ImportError: print( json.dumps( { "success": False, "error": "Pillow is not installed. Install with: pip install Pillow", } ) ) sys.exit(1) except Exception as e: print(json.dumps({"success": False, "error": str(e)})) sys.exit(1) if __name__ == "__main__": main()