Spaces:
Runtime error
Runtime error
| ### | |
| import streamlit as st | |
| import numpy as np | |
| import cv2 | |
| import insightface | |
| from insightface.app import FaceAnalysis | |
| import tempfile | |
| import os | |
| # Initialize face analysis and load model | |
| app = FaceAnalysis(name='buffalo_l') | |
| app.prepare(ctx_id=0, det_size=(640, 640)) | |
| # Load the face swapper model | |
| swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=False, download_zip=False) | |
| def swap_faces_in_video(image, video, progress): | |
| """ | |
| Swaps faces from a source image with faces detected in a video and returns the path to the output video file. | |
| image: Source image (as an array) | |
| video: Path to the input video file | |
| progress: Streamlit progress object | |
| """ | |
| source_faces = app.get(image) | |
| if len(source_faces) == 0: | |
| st.error("No face detected in the source image.") | |
| return None | |
| source_face = source_faces[0] | |
| # Create a temporary file to save the output video | |
| output_path = tempfile.mktemp(suffix='.avi') | |
| # Open the video file | |
| cap = cv2.VideoCapture(video) | |
| # Get video properties for output | |
| frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| # Define the codec and create a VideoWriter object | |
| fourcc = cv2.VideoWriter_fourcc(*'XVID') | |
| out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height)) | |
| for i in range(frame_count): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break # Exit if the video is finished | |
| # Detect faces in the current frame | |
| target_faces = app.get(frame) | |
| # Create a copy of the frame for the result | |
| result_frame = frame.copy() | |
| # Swap faces for each detected face in the video frame | |
| for target_face in target_faces: | |
| result_frame = swapper.get(result_frame, target_face, source_face, paste_back=True) | |
| # Write the result frame to the output video | |
| out.write(result_frame) | |
| # Update progress bar | |
| progress.progress((i + 1) / frame_count) | |
| # Release resources | |
| cap.release() | |
| out.release() | |
| return output_path | |
| # Streamlit UI | |
| st.title("Face Swapper in Video") | |
| st.write("Upload an image and a video to swap faces.") | |
| # File uploader for the source image | |
| image_file = st.file_uploader("Upload Source Image", type=["jpg", "jpeg", "png"]) | |
| # File uploader for the video | |
| video_file = st.file_uploader("Upload Video", type=["mp4", "avi"]) | |
| if st.button("Swap Faces"): | |
| if image_file is not None and video_file is not None: | |
| # Read the source image | |
| source_image = cv2.imdecode(np.frombuffer(image_file.read(), np.uint8), cv2.IMREAD_COLOR) | |
| # Save the uploaded video temporarily | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_video: | |
| tmp_video.write(video_file.read()) | |
| tmp_video_path = tmp_video.name | |
| # Show a spinner and a progress bar while processing | |
| with st.spinner("Processing video..."): | |
| progress_bar = st.progress(0) | |
| output_video_path = swap_faces_in_video(source_image, tmp_video_path, progress_bar) | |
| if output_video_path: | |
| st.success("Face swapping completed!") | |
| # Play the processed video in Streamlit | |
| st.video(output_video_path) | |
| # Provide an option to download the processed video | |
| with open(output_video_path, "rb") as f: | |
| st.download_button( | |
| label="Download Processed Video", | |
| data=f, | |
| file_name="output_swapped_video.avi", | |
| mime="video/x-msvideo" | |
| ) | |
| # Clean up temporary files | |
| os.remove(tmp_video_path) # Clean up temporary video file | |
| # Optionally, keep the output video after displaying | |
| # os.remove(output_video_path) # Uncomment to delete after displaying | |
| else: | |
| st.error("Please upload both an image and a video.") | |