File size: 4,113 Bytes
1bd58ac 3a66796 1bd58ac 3a66796 1bd58ac 3a66796 1bd58ac 3a66796 1bd58ac 3a66796 1bd58ac 3a66796 1bd58ac 3a66796 1bd58ac d5a59cb 1bd58ac d5a59cb 1bd58ac d5a59cb 1bd58ac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
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.")
|