fs-ai / roop /__init__.py
sandhyasinha655's picture
Upload __init__.py
109a39e verified
raw
history blame
1.74 kB
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"):
# Image swap (simplified)
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")