Caveman017/FaceSwapAll-jora-tech

1import cv22import os3import shutil4import time5from SinglePhoto import FaceSwapper6 7def extract_frames(video_path, frames_dir):8    if not os.path.exists(frames_dir):9        os.makedirs(frames_dir)10        last_idx = -111    else:12        # Find the last extracted frame index13        existing = [f for f in os.listdir(frames_dir) if f.startswith("frame_") and f.endswith(".jpg")]14        if existing:15            last_idx = max([int(f.split("_")[1].split(".")[0]) for f in existing])16        else:17            last_idx = -118 19    cap = cv2.VideoCapture(video_path)20    frame_paths = []21    idx = 022    while True:23        ret, frame = cap.read()24        if not ret:25            break26        frame_path = os.path.join(frames_dir, f"frame_{idx:05d}.jpg")27        if idx > last_idx:28            cv2.imwrite(frame_path, frame)29        frame_paths.append(frame_path)30        idx += 131    cap.release()32    return frame_paths33 34def frames_to_video(frames_dir, output_video_path, fps):35    frames = sorted([os.path.join(frames_dir, f) for f in os.listdir(frames_dir) if f.endswith('.jpg')])36    if not frames:37        print("No frames found in directory.")38        return39    first_frame = cv2.imread(frames[0])40    height, width, layers = first_frame.shape41    fourcc = cv2.VideoWriter_fourcc(*'mp4v')42    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))43    for frame_path in frames:44        frame = cv2.imread(frame_path)45        out.write(frame)46    out.release()47 48def main():49    # Use files from VideoSwapping folder50    video_path = os.path.join("VideoSwapping", "data_dst.mp4")51    source_image_path = os.path.join("VideoSwapping", "data_src.jpg")52    frames_dir = os.path.join("VideoSwapping", "video_frames")53    swapped_dir = os.path.join("VideoSwapping", "swapped_frames")54    output_video_path = os.path.join("VideoSwapping", "output_swapped_video.mp4")55    source_face_idx = 156 57    # Ask user for target_face_idx58    while True:59        try:60            user_input = input("Enter target_face_idx (default is 1): ").strip()61            if user_input == "":62                dest_face_idx = 163                break64            dest_face_idx = int(user_input)65            break66        except ValueError:67            print("Invalid input. Please enter an integer value.")68 69    print("Choose an option:")70    print("1. Extract frames only")71    print("2. Face swap only (requires extracted frames)")72    print("3. Both extract frames and face swap")73    choice = input("Enter 1, 2, or 3: ").strip()74 75    frame_paths = []76    if choice == "1" or choice == "3":77        print("Extracting frames from video (resuming if needed)...")78        frame_paths = extract_frames(video_path, frames_dir)79        print(f"Extracted {len(frame_paths)} frames to {frames_dir}.")80        if choice == "1":81            return82 83    if choice == "2":84        # If only face swap, ensure frames are present85        if not os.path.exists(frames_dir):86            print("Frames directory does not exist. Please extract frames first.")87            return88        frame_paths = sorted([os.path.join(frames_dir, f) for f in os.listdir(frames_dir) if f.endswith('.jpg')])89 90    if choice == "2" or choice == "3":91        # Prepare output directory92        if not os.path.exists(swapped_dir):93            os.makedirs(swapped_dir)94 95        # Initialize face swapper96        swapper = FaceSwapper()97 98        # Swap faces on each frame99        print("Swapping faces on frames...")100        start_time = time.time()101        for idx, frame_path in enumerate(frame_paths):102            frame_name = os.path.basename(frame_path)103            # Replace 'frame_' with 'swapped_' in the filename104            if frame_name.startswith("frame_"):105                swapped_name = "swapped_" + frame_name[len("frame_"):]106            else:107                swapped_name = "swapped_" + frame_name108            out_path = os.path.join(swapped_dir, swapped_name)109            if os.path.exists(out_path):110                # Skip already swapped frames111                print(f"Frame {idx+1}/{len(frame_paths)} already swapped, skipping...", end='\r')112                continue113            try:114                try:115                    swapped = swapper.swap_faces(116                        source_path=source_image_path,117                        source_face_idx=source_face_idx,118                        target_path=frame_path,119                        target_face_idx=dest_face_idx120                    )121                except ValueError as ve:122                    if "Target image contains" in str(ve):123                        print(f"\nFrame {idx}: Target face idx {dest_face_idx} not found, trying with idx 1.",end='\r')124                        swapped = swapper.swap_faces(125                            source_path=source_image_path,126                            source_face_idx=source_face_idx,127                            target_path=frame_path,128                            target_face_idx=1129                        )130                    else:131                        raise ve132                cv2.imwrite(out_path, swapped)133            except Exception as e:134                print(f"\nFrame {idx}: {e}")135                # Optionally, copy the original frame if swap fails136                cv2.imwrite(out_path, cv2.imread(frame_path))137            # Estimate time left138            elapsed = time.time() - start_time139            avg_time = elapsed / (idx + 1)140            remaining = avg_time * (len(frame_paths) - (idx + 1))141            mins, secs = divmod(int(remaining), 60)142            print(f"Swapping frame {idx+1}/{len(frame_paths)} | Est. time left: {mins:02d}:{secs:02d}", end='\r')143        print()  # Move to the next line after the loop144 145        # Get FPS from original video146        cap = cv2.VideoCapture(video_path)147        fps = cap.get(cv2.CAP_PROP_FPS)148        cap.release()149 150        # Combine swapped frames into video151        print("Combining swapped frames into video...")152        frames_to_video(swapped_dir, output_video_path, fps)153        print(f"Done! Output video saved as {output_video_path}")154 155        # Ask user if they want to keep the extracted frames and swapped images156        answer = input("Do you want to keep the extracted frames and swapped images? (y/n): ").strip().lower()157        if answer == 'n':158            try:159                shutil.rmtree(frames_dir)160                shutil.rmtree(swapped_dir)161                print("Temporary folders deleted.")162            except Exception as e:163                print(f"Error deleting folders: {e}")164        else:165            print("Temporary folders kept.")166 167if __name__ == "__main__":168    main()