"""Red-eye removal using MediaPipe Face Mesh.""" 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) def main(): input_path = sys.argv[1] output_path = sys.argv[2] settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {} sensitivity = settings.get("sensitivity", 50) strength = settings.get("strength", 70) out_format = settings.get("format", "original") quality = settings.get("quality", 90) # Map sensitivity (0-100) to LAB "a" channel threshold. # Higher sensitivity = lower threshold = more pixels flagged as red. threshold = 170 - (sensitivity / 100) * 50 # Map strength (0-100) to darken factor. # Higher strength = darker correction. darken_factor = 1.0 - (strength / 100) * 0.7 try: emit_progress(10, "Preparing image") from PIL import Image img = Image.open(input_path).convert("RGB") width, height = img.size # Determine output format if out_format == "original": ext = os.path.splitext(input_path)[1].lower() if ext in (".heic", ".heif"): save_format = "PNG" if not output_path.lower().endswith(".png"): output_path = os.path.splitext(output_path)[0] + ".png" elif ext in (".jpg", ".jpeg"): save_format = "JPEG" elif ext == ".webp": save_format = "WEBP" else: save_format = "PNG" elif out_format == "jpeg": save_format = "JPEG" if not output_path.lower().endswith((".jpg", ".jpeg")): output_path = os.path.splitext(output_path)[0] + ".jpg" elif out_format == "webp": save_format = "WEBP" if not output_path.lower().endswith(".webp"): output_path = os.path.splitext(output_path)[0] + ".webp" else: save_format = "PNG" if not output_path.lower().endswith(".png"): output_path = os.path.splitext(output_path)[0] + ".png" format_label = save_format.lower() if format_label == "jpeg": format_label = "jpg" try: import mediapipe as mp import numpy as np import cv2 emit_progress(25, "Detecting faces") img_array = np.array(img) mesh = mp.solutions.face_mesh.FaceMesh( static_image_mode=True, max_num_faces=10, refine_landmarks=True, min_detection_confidence=0.5, ) results = mesh.process(img_array) mesh.close() faces_detected = 0 eyes_corrected = 0 if results.multi_face_landmarks: faces_detected = len(results.multi_face_landmarks) emit_progress(50, "Analyzing eyes") # Iris landmark indices right_iris = [468, 469, 470, 471, 472] # 468 = center left_iris = [473, 474, 475, 476, 477] # 473 = center if faces_detected > 0: all_eyes = [] for face_landmarks in results.multi_face_landmarks: landmarks = face_landmarks.landmark for iris_indices in [right_iris, left_iris]: center_idx = iris_indices[0] contour_indices = iris_indices[1:] cx = int(landmarks[center_idx].x * width) cy = int(landmarks[center_idx].y * height) # Compute radius from contour landmarks radii = [] for idx in contour_indices: px = int(landmarks[idx].x * width) py = int(landmarks[idx].y * height) dist = np.sqrt((px - cx) ** 2 + (py - cy) ** 2) radii.append(dist) radius = np.mean(radii) if radii else 5.0 all_eyes.append((cx, cy, radius)) total_eyes = len(all_eyes) for eye_i, (cx, cy, radius) in enumerate(all_eyes): progress = 50 + int((eye_i + 1) / total_eyes * 40) emit_progress(progress, f"Correcting eye {eye_i + 1} of {total_eyes}") # Padded radius for the circular mask padded_radius = radius * 1.3 r_int = int(np.ceil(padded_radius)) # Bounding box for the ROI x1 = max(0, cx - r_int) y1 = max(0, cy - r_int) x2 = min(width, cx + r_int) y2 = min(height, cy + r_int) if x2 <= x1 or y2 <= y1: continue # Create circular mask in ROI space roi_h = y2 - y1 roi_w = x2 - x1 yy, xx = np.ogrid[:roi_h, :roi_w] local_cx = cx - x1 local_cy = cy - y1 circle_mask = ((xx - local_cx) ** 2 + (yy - local_cy) ** 2) <= (padded_radius ** 2) # Extract ROI roi = img_array[y1:y2, x1:x2].copy() # Convert to LAB roi_lab = cv2.cvtColor(roi, cv2.COLOR_RGB2LAB).astype(np.float64) L_chan = roi_lab[:, :, 0] a_chan = roi_lab[:, :, 1] # LAB red detection: a > threshold AND 50 < L < 220 lab_red = (a_chan > threshold) & (L_chan > 50) & (L_chan < 220) # HSV saturation check roi_hsv = cv2.cvtColor(roi, cv2.COLOR_RGB2HSV) S_chan = roi_hsv[:, :, 1] hsv_saturated = S_chan > 60 # Intersection: LAB-red AND HSV-saturated AND inside circle red_mask = lab_red & hsv_saturated & circle_mask # Morphological cleanup kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)) red_mask_u8 = red_mask.astype(np.uint8) * 255 red_mask_u8 = cv2.morphologyEx(red_mask_u8, cv2.MORPH_CLOSE, kernel) red_mask_u8 = cv2.morphologyEx(red_mask_u8, cv2.MORPH_OPEN, kernel) red_pixel_count = np.count_nonzero(red_mask_u8) if red_pixel_count < 3: continue # Correct red pixels in LAB space corrected_lab = roi_lab.copy() mask_bool = red_mask_u8 > 0 corrected_lab[:, :, 0][mask_bool] = corrected_lab[:, :, 0][mask_bool] * darken_factor corrected_lab[:, :, 1][mask_bool] = 128 # neutral a corrected_lab[:, :, 2][mask_bool] = 128 # neutral b corrected_lab = np.clip(corrected_lab, 0, 255).astype(np.uint8) corrected_rgb = cv2.cvtColor(corrected_lab, cv2.COLOR_LAB2RGB) # Soft mask for blending (Gaussian blur on mask edges) soft_mask = cv2.GaussianBlur( red_mask_u8.astype(np.float32), (5, 5), 1.5 ) soft_mask = soft_mask / 255.0 soft_mask = soft_mask[:, :, np.newaxis] # Alpha blend corrected with original blended = (corrected_rgb.astype(np.float32) * soft_mask + roi.astype(np.float32) * (1.0 - soft_mask)) blended = np.clip(blended, 0, 255).astype(np.uint8) img_array[y1:y2, x1:x2] = blended eyes_corrected += 1 # Update the PIL image from the corrected array img = Image.fromarray(img_array) emit_progress(95, "Saving result") save_kwargs = {} if save_format == "JPEG": save_kwargs["quality"] = quality elif save_format == "WEBP": save_kwargs["quality"] = quality img.save(output_path, format=save_format, **save_kwargs) print( json.dumps( { "success": True, "facesDetected": faces_detected, "eyesCorrected": eyes_corrected, "width": width, "height": height, "format": format_label, "output_path": output_path, } ) ) except ImportError: print( json.dumps( { "success": False, "error": "Red-eye removal requires MediaPipe, NumPy, and OpenCV. Install with: pip install mediapipe numpy opencv-python", } ) ) 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()