SinglePhoto.py

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()