Caveman017/FaceSwapAll-jora-tech

1import cv22import insightface3from insightface.app import FaceAnalysis4import os5 6class FaceSwapper:7    def __init__(self):8        self.app = FaceAnalysis(name='buffalo_l')9        self.app.prepare(ctx_id=0, det_size=(640, 640))10        self.swapper = insightface.model_zoo.get_model(11            'inswapper_128.onnx', download=True, download_zip=True12        )13 14    def swap_faces(self, source_path, source_face_idx, target_path, target_face_idx):15        source_img = cv2.imread(source_path)16        target_img = cv2.imread(target_path)17 18        if source_img is None or target_img is None:19            raise ValueError("Could not read one or both images")20 21        source_faces = self.app.get(source_img)22        target_faces = self.app.get(target_img)23 24        source_faces = sorted(source_faces, key=lambda x: x.bbox[0])25        target_faces = sorted(target_faces, key=lambda x: x.bbox[0])26 27        if len(source_faces) < source_face_idx or source_face_idx < 1:28            raise ValueError(f"Source image contains {len(source_faces)} faces, but requested face {source_face_idx}")29        if len(target_faces) < target_face_idx or target_face_idx < 1:30            raise ValueError(f"Target image contains {len(target_faces)} faces, but requested face {target_face_idx}")31 32        source_face = source_faces[source_face_idx - 1]33        target_face = target_faces[target_face_idx - 1]34 35        result = self.swapper.get(target_img, target_face, source_face, paste_back=True)36        return result37 38    def count_faces(self, img_path):39        """40        Counts the number of faces in the given image file.41        """42        img = cv2.imread(img_path)43        # Use your face detector here. For example, with OpenCV's Haar cascade:44        face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")45        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)46        faces = face_cascade.detectMultiScale(gray, 1.1, 4)47        return len(faces)48 49def main():50    # Paths relative to root51    source_path = os.path.join("SinglePhoto", "data_src.jpg")52    target_path = os.path.join("SinglePhoto", "data_dst.jpg")53    output_dir = os.path.join("SinglePhoto", "output")54    if not os.path.exists(output_dir):55        os.makedirs(output_dir)56 57    swapper = FaceSwapper()58 59    try:60        # Ask user for target_face_idx, default to 1 if no input or invalid input61        try:62            user_input = input("Enter the target face index (starting from 1, default is 1): ")63            target_face_idx = int(user_input) if user_input.strip() else 164            if target_face_idx < 1:65                print("Invalid index. Using default value 1.")66                target_face_idx = 167        except ValueError:68            print("Invalid input. Using default value 1.")69            target_face_idx = 170 71        try:72            result = swapper.swap_faces(73                source_path=source_path,74                source_face_idx=1,75                target_path=target_path,76                target_face_idx=target_face_idx77            )78        except ValueError as ve:79            if "Target image contains" in str(ve):80                print(f"Target face idx {target_face_idx} not found, trying with idx 1.")81                result = swapper.swap_faces(82                    source_path=source_path,83                    source_face_idx=1,84                    target_path=target_path,85                    target_face_idx=186                )87            else:88                raise ve89        output_path = os.path.join(output_dir, "swapped_face.jpg")90        cv2.imwrite(output_path, result)91        print(f"Face swap completed successfully. Result saved to: {output_path}")92    except Exception as e:93        print(f"Error occurred: {str(e)}")94 95if __name__ == "__main__":96    main()