diff --git a/roop/core.py b/roop/core.py index 3ebadb5df..d7f33a824 100644 --- a/roop/core.py +++ b/roop/core.py @@ -12,7 +12,7 @@ from pathlib import Path import tkinter as tk from tkinter import filedialog -# from opennsfw2 import predict_video_frames, predict_image +from opennsfw2 import predict_video_frames, predict_image from tkinter.filedialog import asksaveasfilename import webbrowser import psutil @@ -200,14 +200,14 @@ def start(): if not args['output_file']: args['output_file'] = rreplace(target_path, "/", "/swapped-", 1) if "/" in target_path else "swapped-" + target_path if is_img(target_path): - # if predict_image(target_path) > 0.85: - # quit() + if predict_image(target_path) > 0.85: + quit() process_img(args['source_img'], target_path, os.path.splitext(args['output_file'])[0] + f"_{i+1}.png") status("swap successful!") return - # seconds, probabilities = predict_video_frames(video_path=args['target_path'], frame_interval=100) - # if any(probability > 0.85 for probability in probabilities): - # quit() + seconds, probabilities = predict_video_frames(video_path=args['target_path'], frame_interval=100) + if any(probability > 0.85 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 if os.path.dirname(target_path) else video_name @@ -289,4 +289,4 @@ def run(): 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) - window.mainloop() \ No newline at end of file + window.mainloop()