|
import os |
|
import shutil |
|
import cv2 |
|
from moviepy.editor import VideoFileClip |
|
from PIL import Image |
|
from roop.core import run |
|
from tqdm import tqdm |
|
|
|
def run_face_swap(source_img, target_file): |
|
filename = target_file.name |
|
|
|
if filename.endswith(".jpg") or filename.endswith(".png"): |
|
|
|
os.makedirs("output", exist_ok=True) |
|
output_path = "output/result.jpg" |
|
os.system(f"python run.py --source {source_img} --target {filename} --output {output_path}") |
|
return output_path |
|
|
|
elif filename.endswith(".mp4"): |
|
os.makedirs("frames_output", exist_ok=True) |
|
os.makedirs("frames_input", exist_ok=True) |
|
clip = VideoFileClip(filename) |
|
fps = clip.fps |
|
for i, frame in enumerate(clip.iter_frames()): |
|
Image.fromarray(frame).save(f"frames_input/frame_{i:05d}.jpg") |
|
import sys |
|
sys.argv = ["run.py", "--source", source_img, "--target", "", "--output", ""] |
|
for frame_file in sorted(os.listdir("frames_input")): |
|
input_path = os.path.join("frames_input", frame_file) |
|
output_path = os.path.join("frames_output", frame_file) |
|
sys.argv[4] = input_path |
|
sys.argv[6] = output_path |
|
run() |
|
frame_files = sorted(os.listdir("frames_output")) |
|
frame_images = [cv2.imread(f"frames_output/{f}") for f in frame_files] |
|
h, w, _ = frame_images[0].shape |
|
out_path = "output/result.mp4" |
|
writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) |
|
for img in frame_images: |
|
writer.write(img) |
|
writer.release() |
|
return out_path |
|
|
|
else: |
|
raise ValueError("Unsupported file type") |
|
|