mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
feat(red-eye-removal): SOTA red eye removal with MediaPipe Face Mesh + OpenCV LAB correction (#60)
Uses MediaPipe Face Mesh (refine_landmarks=True) for precise iris localization and OpenCV LAB color space for accurate red-eye detection and luminance-preserving correction. Zero new dependencies - leverages existing MediaPipe + OpenCV stack. - Sensitivity slider (LAB 'a' channel threshold) - Correction strength slider (pupil darkening factor) - Output format selector (Original/PNG/JPEG/WebP) - Before/after preview, progress stages, batch processing - Pipeline support via Controls/Settings split Co-authored-by: stirling-image <stirling-image@users.noreply.github.com>
This commit is contained in:
stirling-image
co-authored by
stirling-image
stirling-image
parent
dfffc0a8cc
commit
9ddeac92b6
@@ -0,0 +1,252 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user