face_utils.py

import cv2
import numpy as np
import mediapipe as mp

# Initialize MediaPipe face modules
mp_face_mesh = mp.solutions.face_mesh
mp_face_detection = mp.solutions.face_detection

# Facial feature indices (MediaPipe 468-point model)
FEATURE_INDICES = {
    'eyes': [33, 133, 362, 263, 155, 154, 153, 145, 144, 163, 7, 173, 157, 158, 159, 160, 161, 246, 
             382, 381, 380, 374, 373, 390, 249, 263, 466, 388, 387, 386, 385, 384, 398],
    'lips': [61, 146, 91, 181, 84, 17, 314, 405, 321, 375, 291, 308, 324, 318, 402, 317, 14, 87, 178, 88, 95],
    'eyebrows': [70, 63, 105, 66, 107, 55, 65, 52, 53, 46, 336, 296, 334, 293, 300, 276, 283, 282, 295, 285],
    'nose': [168, 6, 197, 195, 5, 4, 45, 220, 115, 48, 64, 98, 240]
}

def detect_landmarks(image):
    """Detect facial landmarks using MediaPipe Face Mesh"""
    with mp_face_mesh.FaceMesh(
        static_image_mode=True,
        max_num_faces=1,
        refine_landmarks=True,
        min_detection_confidence=0.5
    ) as face_mesh:
        results = face_mesh.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        
        if not results.multi_face_landmarks:
            return None
            
        landmarks = results.multi_face_landmarks[0].landmark
        h, w = image.shape[:2]
        return [(int(lm.x * w), int(lm.y * h)) for lm in landmarks]

def align_faces(base_img, donor_img):
    """Align donor face to base face using eye positions"""
    base_landmarks = detect_landmarks(base_img)
    donor_landmarks = detect_landmarks(donor_img)
    
    if not base_landmarks or not donor_landmarks:
        return donor_img, donor_landmarks
    
    # Calculate eye centroids
    def get_eye_center(landmarks, indices):
        points = [landmarks[i] for i in indices]
        return np.mean(points, axis=0)
    
    base_left_eye = get_eye_center(base_landmarks, FEATURE_INDICES['eyes'][:6])
    base_right_eye = get_eye_center(base_landmarks, FEATURE_INDICES['eyes'][6:12])
    donor_left_eye = get_eye_center(donor_landmarks, FEATURE_INDICES['eyes'][:6])
    donor_right_eye = get_eye_center(donor_landmarks, FEATURE_INDICES['eyes'][6:12])
    
    # Calculate transformation
    base_eyes_center = np.mean([base_left_eye, base_right_eye], axis=0)
    donor_eyes_center = np.mean([donor_left_eye, donor_right_eye], axis=0)
    
    # Scale based on inter-eye distance
    base_dist = np.linalg.norm(base_left_eye - base_right_eye)
    donor_dist = np.linalg.norm(donor_left_eye - donor_right_eye)
    scale = base_dist / donor_dist
    
    # Rotation angle
    base_angle = np.arctan2(base_right_eye[1] - base_left_eye[1], base_right_eye[0] - base_left_eye[0])
    donor_angle = np.arctan2(donor_right_eye[1] - donor_left_eye[1], donor_right_eye[0] - donor_left_eye[0])
    angle = base_angle - donor_angle
    
    # Rotation matrix
    rot_matrix = cv2.getRotationMatrix2D(tuple(donor_eyes_center), np.degrees(angle), scale)
    
    # Apply transformation
    aligned_donor = cv2.warpAffine(
        donor_img,
        rot_matrix,
        (donor_img.shape[1], donor_img.shape[0]),
        flags=cv2.INTER_CUBIC
    )
    
    # Transform landmarks
    aligned_landmarks = []
    for (x, y) in donor_landmarks:
        aligned_landmarks.append(rot_matrix @ np.array([x, y, 1]))
    
    return aligned_donor, aligned_landmarks

def get_feature_mask(landmarks, feature_type, image_shape):
    """Create mask for specific facial feature"""
    if feature_type not in FEATURE_INDICES:
        raise ValueError(f"Unsupported feature type: {feature_type}")
    
    # Create blank mask
    mask = np.zeros(image_shape[:2], dtype=np.uint8)
    
    # Get feature points
    points = np.array([landmarks[i] for i in FEATURE_INDICES[feature_type]], dtype=np.int32)
    
    # Create convex hull mask
    hull = cv2.convexHull(points)
    cv2.fillConvexPoly(mask, hull, 255)
    
    # Apply Gaussian blur for smooth edges
    mask = cv2.GaussianBlur(mask, (15, 15), 0)
    
    return mask

def blend_features(base_img, donor_img, mask, alpha=0.9):
    """Blend features using seamless cloning and alpha blending"""
    # Find center of mask
    M = cv2.moments(mask)
    cx = int(M["m10"] / M["m00"])
    cy = int(M["m01"] / M["m00"])
    
    # Apply seamless cloning
    blended = cv2.seamlessClone(
        donor_img,
        base_img,
        mask,
        (cx, cy),
        cv2.NORMAL_CLONE
    )
    
    # Apply alpha blending
    mask_float = mask.astype(np.float32) / 255.0 * alpha
    mask_float = np.expand_dims(mask_float, axis=-1)
    result = blended * mask_float + base_img * (1 - mask_float)
    
    return result.astype(np.uint8)

def swap_facial_features(base_img, donor_img, features, blend_alpha=0.9):
    """Main function to swap facial features"""
    # Align donor face to base face
    aligned_donor, aligned_landmarks = align_faces(base_img, donor_img)
    if aligned_landmarks is None:
        return base_img
    
    result = base_img.copy()
    
    # Process each requested feature
    for feature in features:
        feature = feature.lower()
        if feature not in FEATURE_INDICES:
            continue
            
        # Create mask for the feature
        mask = get_feature_mask(aligned_landmarks, feature, base_img.shape)
        
        # Extract feature from donor image
        donor_feature = cv2.bitwise_and(aligned_donor, aligned_donor, mask=mask)
        
        # Blend the feature
        result = blend_features(result, donor_feature, mask, blend_alpha)
    
    return result