SnapOtter/packages/ai/python/red_eye_removal.py

"""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)


# ── Model path for new mp.tasks API ─────────────────────────────────

_MODELS_BASE = os.environ.get("MODELS_PATH", "/opt/models")

_FACE_MESH_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task"
_DOCKER_MODEL_PATH = os.path.join(_MODELS_BASE, "mediapipe", "face_landmarker.task")
_LOCAL_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models")
_LOCAL_MODEL_PATH = os.path.join(_LOCAL_MODEL_DIR, "face_landmarker.task")


def _ensure_face_mesh_model():
    """Resolve face landmarker 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
    from offline_guard import ensure_download_allowed
    ensure_download_allowed("Face landmark model (face_landmarker.task)")
    os.makedirs(_LOCAL_MODEL_DIR, exist_ok=True)
    import urllib.request
    emit_progress(15, "Downloading face mesh model")
    urllib.request.urlretrieve(_FACE_MESH_MODEL_URL, _LOCAL_MODEL_PATH)
    return _LOCAL_MODEL_PATH


def _mesh_with_solutions(img_array, max_faces=50, min_confidence=0.5):
    """FaceMesh using legacy mp.solutions API (mediapipe < 0.10.30).

    Returns list of landmark lists. Each landmark has .x, .y attributes.
    """
    import mediapipe as mp

    mesh = mp.solutions.face_mesh.FaceMesh(
        static_image_mode=True,
        max_num_faces=max_faces,
        refine_landmarks=True,
        min_detection_confidence=min_confidence,
    )
    results = mesh.process(img_array)
    mesh.close()

    if not results.multi_face_landmarks:
        return []

    return [face.landmark for face in results.multi_face_landmarks]


def _mesh_with_tasks(img_array, max_faces=50, min_confidence=0.5):
    """FaceMesh using new mp.tasks API (mediapipe >= 0.10.30).

    Returns list of landmark lists. Each landmark has .x, .y attributes.
    """
    import mediapipe as mp

    model_path = _ensure_face_mesh_model()
    options = mp.tasks.vision.FaceLandmarkerOptions(
        base_options=mp.tasks.BaseOptions(model_asset_path=model_path),
        running_mode=mp.tasks.vision.RunningMode.IMAGE,
        num_faces=max_faces,
        min_face_detection_confidence=min_confidence,
    )
    landmarker = mp.tasks.vision.FaceLandmarker.create_from_options(options)
    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=img_array)
    result = landmarker.detect(mp_image)
    landmarker.close()

    if not result.face_landmarks:
        return []

    return result.face_landmarks


def _detect_face_mesh(img_array, max_faces=50, min_confidence=0.5):
    """Detect face mesh, trying legacy API first then falling back to tasks API."""
    try:
        return _mesh_with_solutions(img_array, max_faces, min_confidence)
    except AttributeError:
        return _mesh_with_tasks(img_array, max_faces, 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 {}

    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 numpy as np
            import cv2

            emit_progress(25, "Detecting faces")

            img_array = np.array(img)

            # Try legacy mp.solutions API first, fall back to mp.tasks
            all_face_landmarks = _detect_face_mesh(img_array)

            faces_detected = len(all_face_landmarks)
            eyes_corrected = 0

            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 landmarks in all_face_landmarks:
                    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 as e:
            msg = str(e)
            hint = "Fix with: apt-get install -y libgl1" if "libGL" in msg else "Install with: pip install mediapipe numpy opencv-python"
            print(
                json.dumps(
                    {
                        "success": False,
                        "error": f"Missing dependency: {msg}. {hint}",
                    }
                )
            )
            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()