Upload selfie_filter.py

selfie_filter.py

@@ -0,0 +1,152 @@
+import os
+import numpy as np
+import cv2
+import landmark_detection
+import gradio as gr
+from mtcnn_facedetection import detect_faces
+
+
+def apply_sunglasses(image, landmarks, sunglasses_img):
+    # If image loading fails or no landmarks, return original image
+    if sunglasses_img is None or not landmarks:
+        return image
+
+    # Create a copy of the image to overlay on
+    result = image.copy()
+
+    # Process each face
+    for face_landmarks in landmarks:
+        # We need at least the eye landmarks
+        if len(face_landmarks) < 5:
+            continue
+
+        # Get eye landmarks
+        left_eye_center = np.mean(face_landmarks[36:42], axis=0).astype(int)
+        right_eye_center = np.mean(face_landmarks[42:48], axis=0).astype(int)
+
+        # Calculate eye distance and angle
+        eye_distance = np.linalg.norm(right_eye_center - left_eye_center)
+        # Negate the angle to correct rotation direction
+        angle = -np.degrees(
+            np.arctan2(
+                right_eye_center[1] - left_eye_center[1],
+                right_eye_center[0] - left_eye_center[0],
+            )
+        )
+
+        # Size for sunglasses based on eye distance
+        width = int(eye_distance * 2.5)
+        height = int(width * sunglasses_img.shape[0] / sunglasses_img.shape[1])
+
+        # Resize sunglasses
+        sunglasses_resized = cv2.resize(sunglasses_img, (width, height))
+
+        # Rotate the sunglasses image
+        center = (width // 2, height // 2)
+        rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
+
+        # Calculate new dimensions after rotation
+        cos = np.abs(rotation_matrix[0, 0])
+        sin = np.abs(rotation_matrix[0, 1])
+        new_width = int((height * sin) + (width * cos))
+        new_height = int((height * cos) + (width * sin))
+
+        # Adjust rotation matrix
+        rotation_matrix[0, 2] += (new_width / 2) - center[0]
+        rotation_matrix[1, 2] += (new_height / 2) - center[1]
+
+        # Perform the rotation
+        rotated_glasses = cv2.warpAffine(
+            sunglasses_resized,
+            rotation_matrix,
+            (new_width, new_height),
+            flags=cv2.INTER_LINEAR,
+            borderMode=cv2.BORDER_CONSTANT,
+            borderValue=(0, 0, 0, 0),
+        )
+
+        # Position the sunglasses
+        eye_center = ((left_eye_center + right_eye_center) // 2).astype(int)
+        x = eye_center[0] - new_width // 2
+        y = eye_center[1] - new_height // 2
+
+        # Create ROI for overlay
+        y1, y2 = max(0, y), min(result.shape[0], y + new_height)
+        x1, x2 = max(0, x), min(result.shape[1], x + new_width)
+
+        # ROI in the glasses image
+        g_y1, g_y2 = max(0, -y), max(0, -y) + (y2 - y1)
+        g_x1, g_x2 = max(0, -x), max(0, -x) + (x2 - x1)
+
+        # Check if we have valid regions
+        if g_y2 <= rotated_glasses.shape[0] and g_x2 <= rotated_glasses.shape[1]:
+            roi = result[y1:y2, x1:x2]
+            glasses_roi = rotated_glasses[g_y1:g_y2, g_x1:g_x2]
+
+            # Apply alpha blending
+            if glasses_roi.shape[2] == 4 and roi.shape[:2] == glasses_roi.shape[:2]:
+                alpha = glasses_roi[:, :, 3] / 255.0
+                for c in range(3):
+                    roi[:, :, c] = (
+                        glasses_roi[:, :, c] * alpha + roi[:, :, c] * (1 - alpha)
+                    ).astype(np.uint8)
+                result[y1:y2, x1:x2] = roi
+
+    return result
+
+
+def do_facial_landmark_recognition(
+    image: np.ndarray, face_boxes: list[landmark_detection.BoundingBox]
+):
+    faces = landmark_detection.get_faces(image, face_boxes)
+    landmarks_batch = landmark_detection.get_landmarks(faces)
+    return landmarks_batch
+
+
+def do_facial_landmark_recognition_with_mtcnn(image: np.ndarray, sunglasses_img):
+    face_boxes = detect_faces(image)
+    landmarks_batch = do_facial_landmark_recognition(image, face_boxes)
+    return apply_sunglasses(image, landmarks_batch, sunglasses_img)
+
+
+def process_video(input_path, sunglasses_img):
+    output_path = os.path.join(
+        os.path.dirname(input_path), "output_" + os.path.basename(input_path)
+    )
+    # Open the input video
+    cap = cv2.VideoCapture(input_path)
+    if not cap.isOpened():
+        gr.Error(f"Error opening input video file: {input_path}")
+        return
+
+    # Get video properties
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    # Create VideoWriter object
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+
+    frame_count = 0
+
+    # Process each frame
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # Process the frame
+        processed_frame = do_facial_landmark_recognition_with_mtcnn(
+            frame, sunglasses_img
+        )
+
+        # Write the frame
+        out.write(processed_frame)
+
+    # Release resources
+    cap.release()
+    out.release()
+    gr.Info(f"Video processing complete. Output saved to: {output_path}")
+    return output_path