mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
9ddeac92b6
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>
253 lines
9.4 KiB
Python
253 lines
9.4 KiB
Python
"""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()
|