debatable/InsightFace-Face_Swapper-on_video

1###2import streamlit as st3import numpy as np4import cv25import insightface6from insightface.app import FaceAnalysis7import tempfile8import os9 10# Initialize face analysis and load model11app = FaceAnalysis(name='buffalo_l')12app.prepare(ctx_id=0, det_size=(640, 640))13 14# Load the face swapper model15swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=False, download_zip=False)16 17def swap_faces_in_video(image, video, progress):18    """19    Swaps faces from a source image with faces detected in a video and returns the path to the output video file.20    21    image: Source image (as an array)22    video: Path to the input video file23    progress: Streamlit progress object24    """25    source_faces = app.get(image)26 27    if len(source_faces) == 0:28        st.error("No face detected in the source image.")29        return None30 31    source_face = source_faces[0]32 33    # Create a temporary file to save the output video34    output_path = tempfile.mktemp(suffix='.avi')35 36    # Open the video file37    cap = cv2.VideoCapture(video)38 39    # Get video properties for output40    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))41    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))42    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))43    fps = cap.get(cv2.CAP_PROP_FPS)44 45    # Define the codec and create a VideoWriter object46    fourcc = cv2.VideoWriter_fourcc(*'XVID')47    out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))48 49    for i in range(frame_count):50        ret, frame = cap.read()51        if not ret:52            break  # Exit if the video is finished53 54        # Detect faces in the current frame55        target_faces = app.get(frame)56 57        # Create a copy of the frame for the result58        result_frame = frame.copy()59 60        # Swap faces for each detected face in the video frame61        for target_face in target_faces:62            result_frame = swapper.get(result_frame, target_face, source_face, paste_back=True)63 64        # Write the result frame to the output video65        out.write(result_frame)66 67        # Update progress bar68        progress.progress((i + 1) / frame_count)69 70    # Release resources71    cap.release()72    out.release()73    74    return output_path75 76# Streamlit UI77st.title("Face Swapper in Video")78st.write("Upload an image and a video to swap faces.")79 80# File uploader for the source image81image_file = st.file_uploader("Upload Source Image", type=["jpg", "jpeg", "png"])82 83# File uploader for the video84video_file = st.file_uploader("Upload Video", type=["mp4", "avi"])85 86if st.button("Swap Faces"):87    if image_file is not None and video_file is not None:88        # Read the source image89        source_image = cv2.imdecode(np.frombuffer(image_file.read(), np.uint8), cv2.IMREAD_COLOR)90        91        # Save the uploaded video temporarily92        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_video:93            tmp_video.write(video_file.read())94            tmp_video_path = tmp_video.name95 96        # Show a spinner and a progress bar while processing97        with st.spinner("Processing video..."):98            progress_bar = st.progress(0)99            output_video_path = swap_faces_in_video(source_image, tmp_video_path, progress_bar)100 101        if output_video_path:102            st.success("Face swapping completed!")103            # Play the processed video in Streamlit104            st.video(output_video_path)105 106            # Provide an option to download the processed video107            with open(output_video_path, "rb") as f:108                st.download_button(109                    label="Download Processed Video",110                    data=f,111                    file_name="output_swapped_video.avi",112                    mime="video/x-msvideo"113                )114 115            # Clean up temporary files116            os.remove(tmp_video_path)  # Clean up temporary video file117            # Optionally, keep the output video after displaying118            # os.remove(output_video_path)  # Uncomment to delete after displaying119    else:120        st.error("Please upload both an image and a video.")121