Update app.py
Browse files
app.py
CHANGED
@@ -0,0 +1,100 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import streamlit as st
|
2 |
+
import numpy as np
|
3 |
+
import cv2
|
4 |
+
import insightface
|
5 |
+
from insightface.app import FaceAnalysis
|
6 |
+
import tempfile
|
7 |
+
import os
|
8 |
+
|
9 |
+
# Initialize face analysis and load model
|
10 |
+
app = FaceAnalysis(name='buffalo_l')
|
11 |
+
app.prepare(ctx_id=0, det_size=(640, 640))
|
12 |
+
|
13 |
+
# Load the face swapper model
|
14 |
+
swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=False, download_zip=False)
|
15 |
+
|
16 |
+
def swap_faces_in_video(image, video):
|
17 |
+
"""
|
18 |
+
Swaps faces from a source image with faces detected in a video and returns the path to the output video file.
|
19 |
+
|
20 |
+
image: Source image (as an array)
|
21 |
+
video: Path to the input video file
|
22 |
+
"""
|
23 |
+
source_faces = app.get(image)
|
24 |
+
|
25 |
+
if len(source_faces) == 0:
|
26 |
+
st.error("No face detected in the source image.")
|
27 |
+
return None
|
28 |
+
|
29 |
+
source_face = source_faces[0]
|
30 |
+
|
31 |
+
# Create a temporary file to save the output video
|
32 |
+
output_path = tempfile.mktemp(suffix='.avi')
|
33 |
+
|
34 |
+
# Open the video file
|
35 |
+
cap = cv2.VideoCapture(video)
|
36 |
+
|
37 |
+
# Get video properties for output
|
38 |
+
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
39 |
+
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
40 |
+
fps = cap.get(cv2.CAP_PROP_FPS)
|
41 |
+
|
42 |
+
# Define the codec and create a VideoWriter object
|
43 |
+
fourcc = cv2.VideoWriter_fourcc(*'XVID')
|
44 |
+
out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
|
45 |
+
|
46 |
+
while True:
|
47 |
+
ret, frame = cap.read()
|
48 |
+
if not ret:
|
49 |
+
break # Exit if the video is finished
|
50 |
+
|
51 |
+
# Detect faces in the current frame
|
52 |
+
target_faces = app.get(frame)
|
53 |
+
|
54 |
+
# Create a copy of the frame for the result
|
55 |
+
result_frame = frame.copy()
|
56 |
+
|
57 |
+
# Swap faces for each detected face in the video frame
|
58 |
+
for target_face in target_faces:
|
59 |
+
result_frame = swapper.get(result_frame, target_face, source_face, paste_back=True)
|
60 |
+
|
61 |
+
# Write the result frame to the output video
|
62 |
+
out.write(result_frame)
|
63 |
+
|
64 |
+
# Release resources
|
65 |
+
cap.release()
|
66 |
+
out.release()
|
67 |
+
|
68 |
+
return output_path
|
69 |
+
|
70 |
+
# Streamlit UI
|
71 |
+
st.title("Face Swapper in Video")
|
72 |
+
st.write("Upload an image and a video to swap faces.")
|
73 |
+
|
74 |
+
# File uploader for the source image
|
75 |
+
image_file = st.file_uploader("Upload Source Image", type=["jpg", "jpeg", "png"])
|
76 |
+
|
77 |
+
# File uploader for the video
|
78 |
+
video_file = st.file_uploader("Upload Video", type=["mp4", "avi"])
|
79 |
+
|
80 |
+
if st.button("Swap Faces"):
|
81 |
+
if image_file is not None and video_file is not None:
|
82 |
+
# Read the source image
|
83 |
+
source_image = cv2.imdecode(np.frombuffer(image_file.read(), np.uint8), cv2.IMREAD_COLOR)
|
84 |
+
|
85 |
+
# Save the uploaded video temporarily
|
86 |
+
with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_video:
|
87 |
+
tmp_video.write(video_file.read())
|
88 |
+
tmp_video_path = tmp_video.name
|
89 |
+
|
90 |
+
# Show a spinner while processing
|
91 |
+
with st.spinner("Processing video..."):
|
92 |
+
output_video_path = swap_faces_in_video(source_image, tmp_video_path)
|
93 |
+
|
94 |
+
if output_video_path:
|
95 |
+
st.success("Face swapping completed!")
|
96 |
+
st.video(output_video_path)
|
97 |
+
os.remove(tmp_video_path) # Clean up temporary video file
|
98 |
+
os.remove(output_video_path) # Clean up output video file if needed
|
99 |
+
else:
|
100 |
+
st.error("Please upload both an image and a video.")
|