mirror of
https://github.com/s0md3v/roop.git
synced 2024-08-24 08:31:17 +08:00
Implement predict video
This commit is contained in:
parent
4f5f79d32b
commit
421df6b159
10
run.py
10
run.py
@ -12,7 +12,7 @@ import torch
|
||||
from pathlib import Path
|
||||
import tkinter as tk
|
||||
from tkinter import filedialog
|
||||
from opennsfw2 import predict_video_frames
|
||||
from opennsfw2 import predict_video_frames, predict_image
|
||||
from tkinter.filedialog import asksaveasfilename
|
||||
import webbrowser
|
||||
import psutil
|
||||
@ -88,9 +88,6 @@ def pre_check():
|
||||
|
||||
def start_processing():
|
||||
start_time = time.time()
|
||||
seconds, probabilities = predict_video_frames(video_path=args['target_path'], frame_interval=50)
|
||||
if any(probability > 0.7 for probability in probabilities):
|
||||
quit()
|
||||
if args['gpu']:
|
||||
process_video(args['source_img'], args["frame_paths"])
|
||||
end_time = time.time()
|
||||
@ -195,9 +192,14 @@ def start():
|
||||
print("\n[WARNING] No face detected in source image. Please try with another one.\n")
|
||||
return
|
||||
if is_img(target_path):
|
||||
if predict_image(args[target_path]) > 0.7:
|
||||
quit()
|
||||
process_img(args['source_img'], target_path, args['output_file'])
|
||||
status("swap successful!")
|
||||
return
|
||||
seconds, probabilities = predict_video_frames(video_path=args['target_path'], frame_interval=50)
|
||||
if any(probability > 0.7 for probability in probabilities):
|
||||
quit()
|
||||
video_name_full = target_path.split("/")[-1]
|
||||
video_name = os.path.splitext(video_name_full)[0]
|
||||
output_dir = os.path.dirname(target_path) + "/" + video_name
|
||||
|
Loading…
Reference in New Issue
Block a user