import cv2 import insightface from insightface.app import FaceAnalysis import os import numpy as np class FaceSwapper: def __init__(self): # Initialize FaceAnalysis with detection and landmark models self.app = FaceAnalysis(name='buffalo_l') self.app.prepare(ctx_id=0, det_size=(640, 640)) # Initialize the swapper model self.swapper = insightface.model_zoo.get_model( 'inswapper_128.onnx', download=True, download_zip=True ) def transplant_hair(self, src_img, dst_img, src_face, dst_face): """ Warps the source hair onto the destination face using Affine Transformation. """ # 1. Get Landmarks (keypoints) src_lm = src_face.kps dst_lm = dst_face.kps # 2. Calculate Affine Transform Matrix to align Source face to Target face # We use the eyes (points 0, 1) and nose (point 2) for alignment src_pts = src_lm[:3] dst_pts = dst_lm[:3] M = cv2.getAffineTransform(src_pts.astype(np.float32), dst_pts.astype(np.float32)) # 3. Warp the entire Source Image to match Target Geometry h, w = dst_img.shape[:2] warped_src = cv2.warpAffine(src_img, M, (w, h), borderMode=cv2.BORDER_REFLECT) # 4. Create a Mask for the Hair (Estimation based on Landmarks) # We assume hair is generally above the eyebrows . # Eyebrow points are indices 17-26 in 68-point models, but insightface buffalo_l uses 5 points usually. # If 5 points: 0,1=eyes, 2=nose, 3,4=mouth corners. # We estimate the forehead/hairline is above the eyes. eye_y = int((dst_lm[0][1] + dst_lm[1][1]) / 2) # Average eye height nose_y = int(dst_lm[2][1]) face_height = nose_y - eye_y # Define the hair region (Everything significantly above the eyes) hair_mask = np.zeros((h, w, 3), dtype=np.float32) # Start the mask slightly above the eyes forehead_line = int(eye_y - (face_height * 0.8)) # Create a soft gradient mask from the forehead up if forehead_line > 0: cv2.rectangle(hair_mask, (0, 0), (w, forehead_line), (1, 1, 1), -1) # Blur the mask heavily to blend the hairline hair_mask = cv2.GaussianBlur(hair_mask, (51, 51), 0) # 5. Blend: (WarpedSource * Mask) + (Target * (1-Mask)) dst_float = dst_img.astype(np.float32) / 255.0 src_float = warped_src.astype(np.float32) / 255.0 final = (src_float * hair_mask) + (dst_float * (1.0 - hair_mask)) final = np.clip(final * 255.0, 0, 255).astype(np.uint8) return final def swap_faces(self, source_path, source_face_idx, target_path, target_face_idx, swap_hair=False): source_img = cv2.imread(source_path) target_img = cv2.imread(target_path) if source_img is None or target_img is None: raise ValueError("Could not read one or both images") # Detect faces source_faces = self.app.get(source_img) target_faces = self.app.get(target_img) # Sort faces from left to right source_faces = sorted(source_faces, key=lambda x: x.bbox[0]) target_faces = sorted(target_faces, key=lambda x: x.bbox[0]) if len(source_faces) < source_face_idx or source_face_idx < 1: raise ValueError(f"Source image contains {len(source_faces)} faces, but requested face {source_face_idx}") if len(target_faces) < target_face_idx or target_face_idx < 1: raise ValueError(f"Target image contains {len(target_faces)} faces, but requested face {target_face_idx}") source_face = source_faces[source_face_idx - 1] target_face = target_faces[target_face_idx - 1] # Step 1: Standard Face Swap (Inswapper) result = self.swapper.get(target_img, target_face, source_face, paste_back=True) # Step 2: Optional Hair Transplant (The new logic) if swap_hair: try: result = self.transplant_hair(source_img, result, source_face, target_face) except Exception as e: print(f"Hair swap failed (fallback to standard swap): {e}") # If hair swap fails, we just return the face swap result pass return result def count_faces(self, img_path): """ Counts the number of faces in the given image file. """ img = cv2.imread(img_path) # Use your face detector here. For example, with OpenCV's Haar cascade: face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml") gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) faces = face_cascade.detectMultiScale(gray, 1.1, 4) return len(faces) def main(): # Paths relative to root source_path = os.path.join("SinglePhoto", "data_src.jpg") target_path = os.path.join("SinglePhoto", "data_dst.jpg") output_dir = os.path.join("SinglePhoto", "output") if not os.path.exists(output_dir): os.makedirs(output_dir) swapper = FaceSwapper() try: # Ask user for target_face_idx, default to 1 if no input or invalid input try: user_input = input("Enter the target face index (starting from 1, default is 1): ") target_face_idx = int(user_input) if user_input.strip() else 1 if target_face_idx < 1: print("Invalid index. Using default value 1.") target_face_idx = 1 except ValueError: print("Invalid input. Using default value 1.") target_face_idx = 1 try: # Default swap_hair to False in CLI mode, or True if you want to test it result = swapper.swap_faces( source_path=source_path, source_face_idx=1, target_path=target_path, target_face_idx=target_face_idx, swap_hair=True # Enabled for testing ) except ValueError as ve: if "Target image contains" in str(ve): print(f"Target face idx {target_face_idx} not found, trying with idx 1.") result = swapper.swap_faces( source_path=source_path, source_face_idx=1, target_path=target_path, target_face_idx=1, swap_hair=True ) else: raise ve output_path = os.path.join(output_dir, "swapped_face.jpg") cv2.imwrite(output_path, result) print(f"Face swap completed successfully. Result saved to: {output_path}") except Exception as e: print(f"Error occurred: {str(e)}") if __name__ == "__main__": main()