mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
SnapOtter
7c70c60b9e
Proxmox LXC and other headless Linux installs lack libgl1, causing all AI tools to show a misleading "install opencv-python-headless" error even though the pip package is already installed. Detect the libGL ImportError and suggest `apt-get install -y libgl1` instead.
323 lines
12 KiB
Python
323 lines
12 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)
|
|
|
|
|
|
# ── 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
|
|
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()
|