From 5dfd1c0ced39ad67ea7323df0459791e448b824e Mon Sep 17 00:00:00 2001 From: KRSHH <136873090+KRSHH@users.noreply.github.com> Date: Thu, 30 Jan 2025 19:20:58 +0530 Subject: [PATCH] Eye Mask --- modules/globals.py | 4 +- modules/processors/frame/face_swapper.py | 252 ++++++++++++++++++++++- modules/ui.py | 49 +++-- 3 files changed, 288 insertions(+), 17 deletions(-) diff --git a/modules/globals.py b/modules/globals.py index 9bd9f4c..add0916 100644 --- a/modules/globals.py +++ b/modules/globals.py @@ -21,7 +21,7 @@ keep_audio = True keep_frames = False many_faces = False map_faces = False -color_correction = False # New global variable for color correction toggle +color_correction = False nsfw_filter = False video_encoder = None video_quality = None @@ -41,5 +41,7 @@ show_mouth_mask_box = False mask_feather_ratio = 8 mask_down_size = 0.50 mask_size = 1 +eyes_mask = False +show_eyes_mask_box = False use_fake_face = False fake_face_path = None diff --git a/modules/processors/frame/face_swapper.py b/modules/processors/frame/face_swapper.py index c188393..cd6f926 100644 --- a/modules/processors/frame/face_swapper.py +++ b/modules/processors/frame/face_swapper.py @@ -74,10 +74,10 @@ def swap_face(source_face: Face, target_face: Face, temp_frame: Frame) -> Frame: temp_frame, target_face, source_face, paste_back=True ) - if modules.globals.mouth_mask: - # Create a mask for the target face - face_mask = create_face_mask(target_face, temp_frame) + # Create face mask for both mouth and eyes masking + face_mask = create_face_mask(target_face, temp_frame) + if modules.globals.mouth_mask: # Create the mouth mask mouth_mask, mouth_cutout, mouth_box, lower_lip_polygon = ( create_lower_mouth_mask(target_face, temp_frame) @@ -94,6 +94,23 @@ def swap_face(source_face: Face, target_face: Face, temp_frame: Frame) -> Frame: swapped_frame, target_face, mouth_mask_data ) + if modules.globals.eyes_mask: + # Create the eyes mask + eyes_mask, eyes_cutout, eyes_box, eyes_polygon = ( + create_eyes_mask(target_face, temp_frame) + ) + + # Apply the eyes area + swapped_frame = apply_eyes_area( + swapped_frame, eyes_cutout, eyes_box, face_mask, eyes_polygon + ) + + if modules.globals.show_eyes_mask_box: + eyes_mask_data = (eyes_mask, eyes_cutout, eyes_box, eyes_polygon) + swapped_frame = draw_eyes_mask_visualization( + swapped_frame, target_face, eyes_mask_data + ) + return swapped_frame @@ -613,3 +630,232 @@ def apply_color_transfer(source, target): source = (source - source_mean) * (target_std / source_std) + target_mean return cv2.cvtColor(np.clip(source, 0, 255).astype("uint8"), cv2.COLOR_LAB2BGR) + + +def create_eyes_mask(face: Face, frame: Frame) -> (np.ndarray, np.ndarray, tuple, np.ndarray): + mask = np.zeros(frame.shape[:2], dtype=np.uint8) + eyes_cutout = None + landmarks = face.landmark_2d_106 + if landmarks is not None: + # Left eye landmarks (87-96) and right eye landmarks (33-42) + left_eye = landmarks[87:96] + right_eye = landmarks[33:42] + + # Calculate centers and dimensions for each eye + left_eye_center = np.mean(left_eye, axis=0).astype(np.int32) + right_eye_center = np.mean(right_eye, axis=0).astype(np.int32) + + # Calculate eye dimensions + def get_eye_dimensions(eye_points): + x_coords = eye_points[:, 0] + y_coords = eye_points[:, 1] + width = int((np.max(x_coords) - np.min(x_coords)) * (1 + modules.globals.mask_down_size)) + height = int((np.max(y_coords) - np.min(y_coords)) * (1 + modules.globals.mask_down_size)) + return width, height + + left_width, left_height = get_eye_dimensions(left_eye) + right_width, right_height = get_eye_dimensions(right_eye) + + # Add extra padding + padding = int(max(left_width, right_width) * 0.2) + + # Calculate bounding box for both eyes + min_x = min(left_eye_center[0] - left_width//2, right_eye_center[0] - right_width//2) - padding + max_x = max(left_eye_center[0] + left_width//2, right_eye_center[0] + right_width//2) + padding + min_y = min(left_eye_center[1] - left_height//2, right_eye_center[1] - right_height//2) - padding + max_y = max(left_eye_center[1] + left_height//2, right_eye_center[1] + right_height//2) + padding + + # Ensure coordinates are within frame bounds + min_x = max(0, min_x) + min_y = max(0, min_y) + max_x = min(frame.shape[1], max_x) + max_y = min(frame.shape[0], max_y) + + # Create mask for the eyes region + mask_roi = np.zeros((max_y - min_y, max_x - min_x), dtype=np.uint8) + + # Draw ellipses for both eyes + left_center = (left_eye_center[0] - min_x, left_eye_center[1] - min_y) + right_center = (right_eye_center[0] - min_x, right_eye_center[1] - min_y) + + # Calculate axes lengths (half of width and height) + left_axes = (left_width//2, left_height//2) + right_axes = (right_width//2, right_height//2) + + # Draw filled ellipses + cv2.ellipse(mask_roi, left_center, left_axes, 0, 0, 360, 255, -1) + cv2.ellipse(mask_roi, right_center, right_axes, 0, 0, 360, 255, -1) + + # Apply Gaussian blur to soften mask edges + mask_roi = cv2.GaussianBlur(mask_roi, (15, 15), 5) + + # Place the mask ROI in the full-sized mask + mask[min_y:max_y, min_x:max_x] = mask_roi + + # Extract the masked area from the frame + eyes_cutout = frame[min_y:max_y, min_x:max_x].copy() + + # Create polygon points for visualization + def create_ellipse_points(center, axes): + t = np.linspace(0, 2*np.pi, 32) + x = center[0] + axes[0] * np.cos(t) + y = center[1] + axes[1] * np.sin(t) + return np.column_stack((x, y)).astype(np.int32) + + # Generate points for both ellipses + left_points = create_ellipse_points((left_eye_center[0], left_eye_center[1]), (left_width//2, left_height//2)) + right_points = create_ellipse_points((right_eye_center[0], right_eye_center[1]), (right_width//2, right_height//2)) + + # Combine points for both eyes + eyes_polygon = np.vstack([left_points, right_points]) + + return mask, eyes_cutout, (min_x, min_y, max_x, max_y), eyes_polygon + + +def apply_eyes_area( + frame: np.ndarray, + eyes_cutout: np.ndarray, + eyes_box: tuple, + face_mask: np.ndarray, + eyes_polygon: np.ndarray, +) -> np.ndarray: + min_x, min_y, max_x, max_y = eyes_box + box_width = max_x - min_x + box_height = max_y - min_y + + if ( + eyes_cutout is None + or box_width is None + or box_height is None + or face_mask is None + or eyes_polygon is None + ): + return frame + + try: + resized_eyes_cutout = cv2.resize(eyes_cutout, (box_width, box_height)) + roi = frame[min_y:max_y, min_x:max_x] + + if roi.shape != resized_eyes_cutout.shape: + resized_eyes_cutout = cv2.resize( + resized_eyes_cutout, (roi.shape[1], roi.shape[0]) + ) + + color_corrected_eyes = apply_color_transfer(resized_eyes_cutout, roi) + + # Create mask for both eyes + polygon_mask = np.zeros(roi.shape[:2], dtype=np.uint8) + + # Split points for left and right eyes + mid_point = len(eyes_polygon) // 2 + left_eye_points = eyes_polygon[:mid_point] - [min_x, min_y] + right_eye_points = eyes_polygon[mid_point:] - [min_x, min_y] + + # Draw filled ellipses using points + left_rect = cv2.minAreaRect(left_eye_points) + right_rect = cv2.minAreaRect(right_eye_points) + + # Convert rect to ellipse parameters + def rect_to_ellipse_params(rect): + center = rect[0] + size = rect[1] + angle = rect[2] + return (int(center[0]), int(center[1])), (int(size[0]/2), int(size[1]/2)), angle + + # Draw filled ellipses + left_params = rect_to_ellipse_params(left_rect) + right_params = rect_to_ellipse_params(right_rect) + cv2.ellipse(polygon_mask, left_params[0], left_params[1], left_params[2], 0, 360, 255, -1) + cv2.ellipse(polygon_mask, right_params[0], right_params[1], right_params[2], 0, 360, 255, -1) + + # Apply feathering + feather_amount = min( + 30, + box_width // modules.globals.mask_feather_ratio, + box_height // modules.globals.mask_feather_ratio, + ) + feathered_mask = cv2.GaussianBlur( + polygon_mask.astype(float), (0, 0), feather_amount + ) + feathered_mask = feathered_mask / feathered_mask.max() + + face_mask_roi = face_mask[min_y:max_y, min_x:max_x] + combined_mask = feathered_mask * (face_mask_roi / 255.0) + + combined_mask = combined_mask[:, :, np.newaxis] + blended = ( + color_corrected_eyes * combined_mask + roi * (1 - combined_mask) + ).astype(np.uint8) + + # Apply face mask to blended result + face_mask_3channel = ( + np.repeat(face_mask_roi[:, :, np.newaxis], 3, axis=2) / 255.0 + ) + final_blend = blended * face_mask_3channel + roi * (1 - face_mask_3channel) + + frame[min_y:max_y, min_x:max_x] = final_blend.astype(np.uint8) + except Exception as e: + pass + + return frame + + +def draw_eyes_mask_visualization( + frame: Frame, face: Face, eyes_mask_data: tuple +) -> Frame: + landmarks = face.landmark_2d_106 + if landmarks is not None and eyes_mask_data is not None: + mask, eyes_cutout, (min_x, min_y, max_x, max_y), eyes_polygon = eyes_mask_data + + vis_frame = frame.copy() + + # Ensure coordinates are within frame bounds + height, width = vis_frame.shape[:2] + min_x, min_y = max(0, min_x), max(0, min_y) + max_x, max_y = min(width, max_x), min(height, max_y) + + # Draw the eyes ellipses + mid_point = len(eyes_polygon) // 2 + left_points = eyes_polygon[:mid_point] + right_points = eyes_polygon[mid_point:] + + try: + # Fit ellipses to points - need at least 5 points + if len(left_points) >= 5 and len(right_points) >= 5: + # Convert points to the correct format for ellipse fitting + left_points = left_points.astype(np.float32) + right_points = right_points.astype(np.float32) + + # Fit ellipses + left_ellipse = cv2.fitEllipse(left_points) + right_ellipse = cv2.fitEllipse(right_points) + + # Draw the ellipses + cv2.ellipse(vis_frame, left_ellipse, (0, 255, 0), 2) + cv2.ellipse(vis_frame, right_ellipse, (0, 255, 0), 2) + except Exception as e: + # If ellipse fitting fails, draw simple rectangles as fallback + left_rect = cv2.boundingRect(left_points) + right_rect = cv2.boundingRect(right_points) + cv2.rectangle(vis_frame, + (left_rect[0], left_rect[1]), + (left_rect[0] + left_rect[2], left_rect[1] + left_rect[3]), + (0, 255, 0), 2) + cv2.rectangle(vis_frame, + (right_rect[0], right_rect[1]), + (right_rect[0] + right_rect[2], right_rect[1] + right_rect[3]), + (0, 255, 0), 2) + + # Add label + cv2.putText( + vis_frame, + "Eyes Mask", + (min_x, min_y - 10), + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + (255, 255, 255), + 1, + ) + + return vis_frame + return frame diff --git a/modules/ui.py b/modules/ui.py index 8242d14..30f1f53 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -3,7 +3,7 @@ import webbrowser import customtkinter as ctk from typing import Callable, Tuple import cv2 -from cv2_enumerate_cameras import enumerate_cameras # Add this import +from cv2_enumerate_cameras import enumerate_cameras from PIL import Image, ImageOps import time import json @@ -252,6 +252,19 @@ def create_root(start: Callable[[], None], destroy: Callable[[], None]) -> ctk.C ) enhancer_switch.place(relx=0.1, rely=0.70) + keep_audio_value = ctk.BooleanVar(value=modules.globals.keep_audio) + keep_audio_switch = ctk.CTkSwitch( + root, + text=_("Keep audio"), + variable=keep_audio_value, + cursor="hand2", + command=lambda: ( + setattr(modules.globals, "keep_audio", keep_audio_value.get()), + save_switch_states(), + ), + ) + keep_audio_switch.place(relx=0.1, rely=0.75) + # Additional Options (Middle Right) mouth_mask_var = ctk.BooleanVar(value=modules.globals.mouth_mask) mouth_mask_switch = ctk.CTkSwitch( @@ -273,20 +286,31 @@ def create_root(start: Callable[[], None], destroy: Callable[[], None]) -> ctk.C modules.globals, "show_mouth_mask_box", show_mouth_mask_box_var.get() ), ) - show_mouth_mask_box_switch.place(relx=0.6, rely=0.6) + show_mouth_mask_box_switch.place(relx=0.6, rely=0.60) - keep_audio_value = ctk.BooleanVar(value=modules.globals.keep_audio) - keep_audio_switch = ctk.CTkSwitch( + # Add eyes mask switch + eyes_mask_var = ctk.BooleanVar(value=modules.globals.eyes_mask) + eyes_mask_switch = ctk.CTkSwitch( root, - text=_("Keep audio"), - variable=keep_audio_value, + text=_("Eyes Mask"), + variable=eyes_mask_var, cursor="hand2", - command=lambda: ( - setattr(modules.globals, "keep_audio", keep_audio_value.get()), - save_switch_states(), + command=lambda: setattr(modules.globals, "eyes_mask", eyes_mask_var.get()), + ) + eyes_mask_switch.place(relx=0.6, rely=0.65) + + # Add show eyes mask box switch + show_eyes_mask_box_var = ctk.BooleanVar(value=modules.globals.show_eyes_mask_box) + show_eyes_mask_box_switch = ctk.CTkSwitch( + root, + text=_("Show Eyes Mask Box"), + variable=show_eyes_mask_box_var, + cursor="hand2", + command=lambda: setattr( + modules.globals, "show_eyes_mask_box", show_eyes_mask_box_var.get() ), ) - keep_audio_switch.place(relx=0.6, rely=0.65) + show_eyes_mask_box_switch.place(relx=0.6, rely=0.70) # Add show FPS switch show_fps_value = ctk.BooleanVar(value=modules.globals.show_fps) @@ -300,7 +324,7 @@ def create_root(start: Callable[[], None], destroy: Callable[[], None]) -> ctk.C save_switch_states(), ), ) - show_fps_switch.place(relx=0.6, rely=0.70) + show_fps_switch.place(relx=0.6, rely=0.75) # Main Control Buttons (Bottom) start_button = ctk.CTkButton( @@ -767,8 +791,7 @@ def update_preview(frame_number: int = 0) -> None: modules.globals.frame_processors ): temp_frame = frame_processor.process_frame( - get_one_face(cv2.imread(modules.globals.source_path)), temp_frame - ) + get_one_face(cv2.imread(modules.globals.source_path)), temp_frame) image = Image.fromarray(cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB)) image = ImageOps.contain( image, (PREVIEW_MAX_WIDTH, PREVIEW_MAX_HEIGHT), Image.LANCZOS