2023-05-30 00:40:02 +02:00
#!/usr/bin/env python3
2023-05-29 19:04:00 +05:30
import sys
2023-05-30 01:28:16 +02:00
import time
2023-05-29 23:35:29 +05:30
import shutil
2023-05-29 19:04:00 +05:30
import core . globals
2023-05-29 23:35:29 +05:30
if not shutil . which ( ' ffmpeg ' ) :
2023-05-30 01:20:40 +05:30
print ( ' ffmpeg is not installed. Read the docs: https://github.com/s0md3v/roop#installation. \n ' * 10 )
2023-05-29 23:57:05 +05:30
quit ( )
2023-05-29 23:10:52 +05:30
if ' --gpu ' not in sys . argv :
core . globals . providers = [ ' CPUExecutionProvider ' ]
2023-05-30 08:11:55 +05:30
elif ' ROCMExecutionProvider ' not in core . globals . providers :
2023-05-30 07:24:10 +05:30
import torch
2023-05-30 08:50:04 +05:30
if not torch . cuda . is_available ( ) :
2023-05-30 08:52:21 +05:30
quit ( " You are using --gpu flag but CUDA isn ' t available or properly installed on your system. " )
2023-05-30 07:24:10 +05:30
2023-05-28 20:19:40 +05:30
import glob
import argparse
import multiprocessing as mp
import os
from pathlib import Path
import tkinter as tk
from tkinter import filedialog
2023-05-29 13:39:44 +05:30
from tkinter . filedialog import asksaveasfilename
2023-05-28 20:19:40 +05:30
from core . processor import process_video , process_img
from core . utils import is_img , detect_fps , set_fps , create_video , add_audio , extract_frames
2023-05-29 17:50:42 +05:30
from core . config import get_face
2023-05-28 20:19:40 +05:30
import webbrowser
import psutil
2023-05-29 17:50:42 +05:30
import cv2
2023-05-30 00:37:43 +05:30
import threading
from PIL import Image , ImageTk
2023-05-28 20:19:40 +05:30
pool = None
args = { }
parser = argparse . ArgumentParser ( )
parser . add_argument ( ' -f ' , ' --face ' , help = ' use this face ' , dest = ' source_img ' )
parser . add_argument ( ' -t ' , ' --target ' , help = ' replace this face ' , dest = ' target_path ' )
2023-05-29 16:12:06 +03:00
parser . add_argument ( ' -o ' , ' --output ' , help = ' save output to this file ' , dest = ' output_file ' )
2023-05-28 20:19:40 +05:30
parser . add_argument ( ' --keep-fps ' , help = ' maintain original fps ' , dest = ' keep_fps ' , action = ' store_true ' , default = False )
2023-05-30 07:00:37 +05:30
parser . add_argument ( ' --gpu ' , help = ' use gpu ' , dest = ' gpu ' , action = ' store_true ' , default = False )
2023-05-28 20:19:40 +05:30
parser . add_argument ( ' --keep-frames ' , help = ' keep frames directory ' , dest = ' keep_frames ' , action = ' store_true ' , default = False )
2023-05-30 20:30:35 +05:30
parser . add_argument ( ' --cores ' , help = ' number of cores to use ' , dest = ' cores_count ' , type = int )
2023-05-28 20:19:40 +05:30
for name , value in vars ( parser . parse_args ( ) ) . items ( ) :
args [ name ] = value
2023-05-30 20:30:35 +05:30
if not args [ ' cores_count ' ] :
args [ ' cores_count ' ] = psutil . cpu_count ( ) - 1
2023-05-29 18:53:10 +05:30
2023-05-29 22:55:47 +05:30
sep = " / "
if os . name == " nt " :
sep = " \\ "
2023-05-28 20:19:40 +05:30
def start_processing ( ) :
2023-05-30 01:28:16 +02:00
start_time = time . time ( )
2023-05-28 20:19:40 +05:30
if args [ ' gpu ' ] :
process_video ( args [ ' source_img ' ] , args [ " frame_paths " ] )
2023-05-30 01:28:42 +02:00
end_time = time . time ( )
print ( flush = True )
2023-05-30 01:28:16 +02:00
print ( f " Processing time: { end_time - start_time : .2f } seconds " , flush = True )
2023-05-28 20:19:40 +05:30
return
frame_paths = args [ " frame_paths " ]
2023-05-30 20:30:35 +05:30
n = len ( frame_paths ) / / ( args [ ' cores_count ' ] )
2023-05-28 20:19:40 +05:30
processes = [ ]
for i in range ( 0 , len ( frame_paths ) , n ) :
p = pool . apply_async ( process_video , args = ( args [ ' source_img ' ] , frame_paths [ i : i + n ] , ) )
processes . append ( p )
for p in processes :
p . get ( )
pool . close ( )
pool . join ( )
2023-05-30 01:28:16 +02:00
end_time = time . time ( )
2023-05-30 01:28:42 +02:00
print ( flush = True )
2023-05-30 01:28:16 +02:00
print ( f " Processing time: { end_time - start_time : .2f } seconds " , flush = True )
2023-05-28 20:19:40 +05:30
2023-05-30 00:37:43 +05:30
def preview_image ( image_path ) :
img = Image . open ( image_path )
2023-05-30 06:40:46 +05:30
img = img . resize ( ( 180 , 180 ) , Image . ANTIALIAS )
2023-05-30 00:37:43 +05:30
photo_img = ImageTk . PhotoImage ( img )
left_frame = tk . Frame ( window )
2023-05-30 06:40:46 +05:30
left_frame . place ( x = 60 , y = 100 )
2023-05-30 00:37:43 +05:30
img_label = tk . Label ( left_frame , image = photo_img )
img_label . image = photo_img
img_label . pack ( )
def preview_video ( video_path ) :
cap = cv2 . VideoCapture ( video_path )
if not cap . isOpened ( ) :
print ( " Error opening video file " )
return
ret , frame = cap . read ( )
if ret :
frame = cv2 . cvtColor ( frame , cv2 . COLOR_BGR2RGB )
img = Image . fromarray ( frame )
2023-05-30 06:40:46 +05:30
img = img . resize ( ( 180 , 180 ) , Image . ANTIALIAS )
2023-05-30 00:37:43 +05:30
photo_img = ImageTk . PhotoImage ( img )
right_frame = tk . Frame ( window )
2023-05-30 06:40:46 +05:30
right_frame . place ( x = 360 , y = 100 )
2023-05-30 00:37:43 +05:30
img_label = tk . Label ( right_frame , image = photo_img )
img_label . image = photo_img
img_label . pack ( )
cap . release ( )
2023-05-28 20:19:40 +05:30
def select_face ( ) :
args [ ' source_img ' ] = filedialog . askopenfilename ( title = " Select a face " )
2023-05-30 00:37:43 +05:30
preview_image ( args [ ' source_img ' ] )
2023-05-28 20:19:40 +05:30
def select_target ( ) :
args [ ' target_path ' ] = filedialog . askopenfilename ( title = " Select a target " )
2023-05-30 00:37:43 +05:30
threading . Thread ( target = preview_video , args = ( args [ ' target_path ' ] , ) ) . start ( )
2023-05-28 20:19:40 +05:30
def toggle_fps_limit ( ) :
args [ ' keep_fps ' ] = limit_fps . get ( ) != True
2023-05-30 04:27:42 +05:30
def toggle_keep_frames ( ) :
args [ ' keep_frames ' ] = keep_frames . get ( ) != True
2023-05-29 13:39:44 +05:30
def save_file ( ) :
2023-05-30 19:07:22 +05:30
filename , ext = ' output.mp4 ' , ' .mp4 '
if is_img ( args [ ' target_path ' ] ) :
filename , ext = ' output.png ' , ' .png '
args [ ' output_file ' ] = asksaveasfilename ( initialfile = filename , defaultextension = ext , filetypes = [ ( " All Files " , " *.* " ) , ( " Videos " , " *.mp4 " ) ] )
2023-05-29 13:39:44 +05:30
2023-05-30 02:28:31 +05:30
def status ( string ) :
2023-05-30 18:53:31 +05:30
if ' cli_mode ' in args :
2023-05-30 05:26:50 +02:00
print ( " Status: " + string )
else :
status_label [ " text " ] = " Status: " + string
window . update ( )
2023-05-30 18:53:31 +05:30
2023-05-30 02:28:31 +05:30
2023-05-28 20:19:40 +05:30
def start ( ) :
2023-05-30 02:28:31 +05:30
print ( " DON ' T WORRY. IT ' S NOT STUCK/CRASHED. \n " * 5 )
2023-05-29 13:01:05 +05:30
if not args [ ' source_img ' ] or not os . path . isfile ( args [ ' source_img ' ] ) :
print ( " \n [WARNING] Please select an image containing a face. " )
return
elif not args [ ' target_path ' ] or not os . path . isfile ( args [ ' target_path ' ] ) :
print ( " \n [WARNING] Please select a video/image to swap face in. " )
return
2023-05-30 19:07:22 +05:30
if not args [ ' output_file ' ] :
args [ ' output_file ' ] = rreplace ( args [ ' target_path ' ] , " / " , " /swapped- " , 1 ) if " / " in target_path else " swapped- " + target_path
2023-05-28 20:19:40 +05:30
global pool
2023-05-30 20:30:35 +05:30
pool = mp . Pool ( args [ ' cores_count ' ] )
2023-05-28 20:19:40 +05:30
target_path = args [ ' target_path ' ]
2023-05-29 17:50:42 +05:30
test_face = get_face ( cv2 . imread ( args [ ' source_img ' ] ) )
if not test_face :
print ( " \n [WARNING] No face detected in source image. Please try with another one. \n " )
return
2023-05-28 20:19:40 +05:30
if is_img ( target_path ) :
2023-05-30 19:07:22 +05:30
process_img ( args [ ' source_img ' ] , target_path , args [ ' output_file ' ] )
2023-05-30 04:27:42 +05:30
status ( " swap successful! " )
2023-05-28 20:19:40 +05:30
return
2023-05-30 19:29:33 +03:00
video_name = os . path . basename ( target_path )
video_name = os . path . splitext ( video_name ) [ 0 ]
output_dir = os . path . join ( os . path . dirname ( target_path ) , video_name )
2023-05-30 22:04:38 +05:30
Path ( output_dir ) . mkdir ( exist_ok = True )
2023-05-30 04:27:42 +05:30
status ( " detecting video ' s FPS... " )
2023-05-28 20:19:40 +05:30
fps = detect_fps ( target_path )
if not args [ ' keep_fps ' ] and fps > 30 :
2023-05-29 21:24:40 +05:30
this_path = output_dir + " / " + video_name + " .mp4 "
2023-05-28 20:19:40 +05:30
set_fps ( target_path , this_path , 30 )
target_path , fps = this_path , 30
else :
shutil . copy ( target_path , output_dir )
2023-05-30 04:27:42 +05:30
status ( " extracting frames... " )
2023-05-28 20:19:40 +05:30
extract_frames ( target_path , output_dir )
args [ ' frame_paths ' ] = tuple ( sorted (
2023-05-29 21:24:40 +05:30
glob . glob ( output_dir + f " /*.png " ) ,
2023-05-29 22:55:47 +05:30
key = lambda x : int ( x . split ( sep ) [ - 1 ] . replace ( " .png " , " " ) )
2023-05-28 20:19:40 +05:30
) )
2023-05-30 04:27:42 +05:30
status ( " swapping in progress... " )
2023-05-28 20:19:40 +05:30
start_processing ( )
2023-05-30 04:27:42 +05:30
status ( " creating video... " )
2023-05-28 20:19:40 +05:30
create_video ( video_name , fps , output_dir )
2023-05-30 04:27:42 +05:30
status ( " adding audio... " )
2023-05-29 17:42:52 +05:30
add_audio ( output_dir , target_path , args [ ' keep_frames ' ] , args [ ' output_file ' ] )
2023-05-29 21:24:40 +05:30
save_path = args [ ' output_file ' ] if args [ ' output_file ' ] else output_dir + " / " + video_name + " .mp4 "
2023-05-29 15:27:52 +05:30
print ( " \n \n Video saved as: " , save_path , " \n \n " )
2023-05-30 04:27:42 +05:30
status ( " swap successful! " )
2023-05-28 20:19:40 +05:30
if __name__ == " __main__ " :
2023-05-30 02:28:31 +05:30
global status_label , window
2023-05-28 20:19:40 +05:30
if args [ ' source_img ' ] :
2023-05-30 18:53:31 +05:30
args [ ' cli_mode ' ] = True
2023-05-28 20:19:40 +05:30
start ( )
quit ( )
window = tk . Tk ( )
2023-05-30 06:40:46 +05:30
window . geometry ( " 600x700 " )
2023-05-28 20:19:40 +05:30
window . title ( " roop " )
2023-05-30 04:27:42 +05:30
window . configure ( bg = " #2d3436 " )
window . resizable ( width = False , height = False )
2023-05-28 20:19:40 +05:30
# Contact information
2023-05-30 04:27:42 +05:30
support_link = tk . Label ( window , text = " Donate to project <3 " , fg = " #fd79a8 " , bg = " #2d3436 " , cursor = " hand2 " , font = ( " Arial " , 8 ) )
support_link . place ( x = 180 , y = 20 , width = 250 , height = 30 )
2023-05-28 20:19:40 +05:30
support_link . bind ( " <Button-1> " , lambda e : webbrowser . open ( " https://github.com/sponsors/s0md3v " ) )
# Select a face button
2023-05-30 04:27:42 +05:30
face_button = tk . Button ( window , text = " Select a face " , command = select_face , bg = " #2d3436 " , fg = " #74b9ff " , highlightthickness = 4 , relief = " flat " , highlightbackground = " #74b9ff " , activebackground = " #74b9ff " , borderwidth = 4 )
2023-05-30 06:40:46 +05:30
face_button . place ( x = 60 , y = 320 , width = 180 , height = 80 )
2023-05-28 20:19:40 +05:30
# Select a target button
2023-05-30 06:40:46 +05:30
target_button = tk . Button ( window , text = " Select a target " , command = select_target , bg = " #2d3436 " , fg = " #74b9ff " , highlightthickness = 4 , relief = " flat " , highlightbackground = " #74b9ff " , activebackground = " #74b9ff " , borderwidth = 4 )
target_button . place ( x = 360 , y = 320 , width = 180 , height = 80 )
2023-05-28 20:19:40 +05:30
# FPS limit checkbox
limit_fps = tk . IntVar ( )
2023-05-30 04:27:42 +05:30
fps_checkbox = tk . Checkbutton ( window , relief = " groove " , activebackground = " #2d3436 " , activeforeground = " #74b9ff " , selectcolor = " black " , text = " Limit FPS to 30 " , fg = " #dfe6e9 " , borderwidth = 0 , highlightthickness = 0 , bg = " #2d3436 " , variable = limit_fps , command = toggle_fps_limit )
2023-05-30 06:40:46 +05:30
fps_checkbox . place ( x = 30 , y = 500 , width = 240 , height = 31 )
2023-05-28 20:19:40 +05:30
fps_checkbox . select ( )
2023-05-30 04:27:42 +05:30
# Keep frames checkbox
keep_frames = tk . IntVar ( )
frames_checkbox = tk . Checkbutton ( window , relief = " groove " , activebackground = " #2d3436 " , activeforeground = " #74b9ff " , selectcolor = " black " , text = " Keep frames dir " , fg = " #dfe6e9 " , borderwidth = 0 , highlightthickness = 0 , bg = " #2d3436 " , variable = keep_frames , command = toggle_keep_frames )
2023-05-30 06:40:46 +05:30
frames_checkbox . place ( x = 37 , y = 450 , width = 240 , height = 31 )
2023-05-28 20:19:40 +05:30
# Start button
2023-05-30 06:40:46 +05:30
start_button = tk . Button ( window , text = " Start " , bg = " #f1c40f " , relief = " flat " , borderwidth = 0 , highlightthickness = 0 , command = lambda : [ save_file ( ) , start ( ) ] )
start_button . place ( x = 240 , y = 560 , width = 120 , height = 49 )
2023-05-30 02:28:31 +05:30
# Status label
2023-05-30 06:40:46 +05:30
status_label = tk . Label ( window , width = 580 , justify = " center " , text = " Status: waiting for input... " , fg = " #2ecc71 " , bg = " #2d3436 " )
status_label . place ( x = 10 , y = 640 , width = 580 , height = 30 )
2023-05-28 20:19:40 +05:30
window . mainloop ( )