File size: 3,989 Bytes
ed6bcfa 1ccca68 ed6bcfa eb130f8 ed6bcfa eb130f8 ed6bcfa 1ccca68 eb130f8 1ccca68 eb130f8 ed6bcfa eb130f8 ed6bcfa 1ccca68 ed6bcfa eb130f8 | 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 121 122 123 124 125 126 127 | import gradio as gr
import insightface
from insightface.app import FaceAnalysis
import numpy as np
from PIL import Image
import cv2
import tempfile
import os
MODEL_PATH = "inswapper_128.onnx"
# --------- Load models once ---------
print("Loading models...")
face_app = FaceAnalysis(name="buffalo_l")
face_app.prepare(ctx_id=0, det_size=(640, 640)) # ctx_id=0 -> CPU on Spaces free tier
swapper = insightface.model_zoo.get_model(MODEL_PATH)
print("Models loaded.")
# --------- IMAGE SWAP ---------
def swap_faces_image(source, target):
if source is None or target is None:
raise gr.Error("Please upload both source and target images.")
# Gradio gives numpy arrays (H, W, 3) in RGB
src = np.array(source)
dst = np.array(target)
src_faces = face_app.get(src)
dst_faces = face_app.get(dst)
if len(src_faces) == 0:
raise gr.Error("No face found in source image.")
if len(dst_faces) == 0:
raise gr.Error("No face found in target image.")
src_face = src_faces[0]
dst_face = dst_faces[0]
result = swapper.get(dst.copy(), dst_face, src_face, paste_back=True)
return Image.fromarray(result)
# --------- VIDEO SWAP ---------
def swap_faces_video(source_image, video_file):
if source_image is None:
raise gr.Error("Please upload a source face image.")
if video_file is None:
raise gr.Error("Please upload a target video (mp4).")
# source_image: numpy RGB
src = np.array(source_image)
src_faces = face_app.get(src)
if len(src_faces) == 0:
raise gr.Error("No face found in source image.")
src_face = src_faces[0]
# video_file is a file-like object; get its path
video_path = video_file.name
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
raise gr.Error("Could not open uploaded video.")
fps = cap.get(cv2.CAP_PROP_FPS) or 25
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
# temp output file
tmp_out = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
tmp_out_path = tmp_out.name
tmp_out.close()
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
writer = cv2.VideoWriter(tmp_out_path, fourcc, fps, (w, h))
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
print(f"[INFO] Total frames: {total_frames}")
frame_idx = 0
while True:
ret, frame = cap.read()
if not ret:
break
frame_idx += 1
if frame_idx % 10 == 0:
print(f"[INFO] Frame {frame_idx}/{total_frames}")
# BGR -> RGB
rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
dst_faces = face_app.get(rgb)
if len(dst_faces) > 0:
dst_face = dst_faces[0]
swapped_rgb = swapper.get(rgb.copy(), dst_face, src_face, paste_back=True)
out_frame = cv2.cvtColor(swapped_rgb, cv2.COLOR_RGB2BGR)
writer.write(out_frame)
else:
writer.write(frame)
cap.release()
writer.release()
print(f"[INFO] Video finished: {tmp_out_path}")
return tmp_out_path
# --------- Gradio UI (Image + Video) ---------
with gr.Blocks() as demo:
gr.Markdown("## InsightFace FaceSwap (Image & Video) on HuggingFace")
with gr.Tab("Image Swap"):
src_img = gr.Image(type="numpy", label="Your Face (Source)")
tgt_img = gr.Image(type="numpy", label="Target Face")
out_img = gr.Image(label="Swapped Result")
btn_img = gr.Button("Swap Image")
btn_img.click(fn=swap_faces_image, inputs=[src_img, tgt_img], outputs=out_img)
with gr.Tab("Video Swap"):
src_vid_img = gr.Image(type="numpy", label="Your Face (Source)")
tgt_vid = gr.File(label="Target Video (mp4)")
out_vid = gr.Video(label="Swapped Video")
btn_vid = gr.Button("Swap Video")
btn_vid.click(fn=swap_faces_video, inputs=[src_vid_img, tgt_vid], outputs=out_vid)
if __name__ == "__main__":
demo.launch()
|