|
|
import os, cv2, tempfile |
|
|
import gradio as gr |
|
|
from moviepy.editor import ImageSequenceClip |
|
|
from insightface.app import FaceAnalysis |
|
|
from insightface.model_zoo import get_model |
|
|
|
|
|
os.environ["XDG_RUNTIME_DIR"] = "/tmp/runtime-dir" |
|
|
|
|
|
|
|
|
app = FaceAnalysis(name="buffalo_l", providers=["CUDAExecutionProvider","CPUExecutionProvider"]) |
|
|
app.prepare(ctx_id=0, det_size=(1024,1024)) |
|
|
|
|
|
onnx_model_path = "inswapper_128.onnx" |
|
|
swap_model = get_model(onnx_model_path, |
|
|
providers=["CUDAExecutionProvider","CPUExecutionProvider"]) |
|
|
|
|
|
def face_swap_api(source_image, target_video): |
|
|
src_img = cv2.imread(source_image) |
|
|
src_faces = app.get(src_img) |
|
|
if len(src_faces) == 0: |
|
|
return None |
|
|
src_face = src_faces[0] |
|
|
|
|
|
frames_dir = tempfile.mkdtemp() |
|
|
cap = cv2.VideoCapture(target_video) |
|
|
fps = cap.get(cv2.CAP_PROP_FPS) |
|
|
frames = [] |
|
|
idx = 0 |
|
|
while cap.isOpened(): |
|
|
ret, frame = cap.read() |
|
|
if not ret: break |
|
|
frame_path = os.path.join(frames_dir, f"frame_{idx:04d}.png") |
|
|
cv2.imwrite(frame_path, frame) |
|
|
frames.append(frame_path) |
|
|
idx += 1 |
|
|
cap.release() |
|
|
|
|
|
swapped_frames = [] |
|
|
for fpath in frames: |
|
|
frame = cv2.imread(fpath) |
|
|
faces = app.get(frame) |
|
|
for face in faces: |
|
|
frame = swap_model.get(frame, face, src_face, paste_back=True) |
|
|
swapped_frames.append(frame) |
|
|
|
|
|
out_path = os.path.join(tempfile.gettempdir(), "swapped_output.mp4") |
|
|
clip = ImageSequenceClip( |
|
|
[cv2.cvtColor(f, cv2.COLOR_BGR2RGB) for f in swapped_frames], fps=fps |
|
|
) |
|
|
clip.write_videofile(out_path, codec="libx264", audio=False, verbose=False) |
|
|
return out_path |
|
|
|
|
|
demo = gr.Interface( |
|
|
fn=face_swap_api, |
|
|
inputs=[gr.Image(type="filepath", label="Source Face"), |
|
|
gr.Video(label="Target Video")], |
|
|
outputs=gr.Video(label="Swapped Output Video"), |
|
|
title="Face Swap API", |
|
|
description="Upload a face photo and a video. Returns a face-swapped MP4." |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo.launch() |
|
|
|
