diff --git a/README.md b/README.md index 3ba3cc7..56b5c25 100644 --- a/README.md +++ b/README.md @@ -38,9 +38,9 @@ options: --keep-fps maintain original fps --keep-frames keep frames directory --max-memory MAX_MEMORY - set max memory + maximum amount of RAM in GB to be used --max-cores CORES_COUNT - set max cpu cores + number of cores to be use for CPU mode ``` Looking for a CLI mode? Using the -f/--face argument will make the program in cli mode. diff --git a/core/utils.py b/core/utils.py index ee2d422..9a1498e 100644 --- a/core/utils.py +++ b/core/utils.py @@ -23,7 +23,7 @@ def detect_fps(input_path): output = os.popen(f'ffprobe -v error -select_streams v -of default=noprint_wrappers=1:nokey=1 -show_entries stream=r_frame_rate "{input_path}"').read() if "/" in output: try: - return int(output.split("/")[0]) // int(output.split("/")[1]), output.removesuffix('\n') + return int(output.split(os.sep)[0]) // int(output.split(os.sep)[1]), output.removesuffix('\n') except: pass return 30, 30 diff --git a/run.py b/run.py old mode 100644 new mode 100755 index 049c5c4..03b7bb9 --- a/run.py +++ b/run.py @@ -8,24 +8,24 @@ import glob import argparse import multiprocessing as mp import os -import random +import torch from pathlib import Path import tkinter as tk from tkinter import filedialog -from opennsfw2 import predict_image as face_check +from opennsfw2 import predict_video_frames, predict_image from tkinter.filedialog import asksaveasfilename -import core.globals -from core.processor import process_video, process_img -from core.utils import is_img, detect_fps, set_fps, create_video, add_audio, extract_frames, rreplace -from core.config import get_face import webbrowser import psutil import cv2 import threading from PIL import Image, ImageTk +import core.globals +from core.processor import process_video, process_img +from core.utils import is_img, detect_fps, set_fps, create_video, add_audio, extract_frames, rreplace +from core.config import get_face -if 'ROCMExecutionProvider' not in core.globals.providers: - import torch +if 'ROCMExecutionProvider' in core.globals.providers: + del torch pool = None args = {} @@ -37,9 +37,9 @@ parser.add_argument('-o', '--output', help='save output to this file', dest='out parser.add_argument('--gpu', help='use gpu', dest='gpu', action='store_true', default=False) parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False) parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False) -parser.add_argument('--max-memory', help='set max memory', type=int) -parser.add_argument('--max-cores', help='number of cores to use', dest='cores_count', type=int, default=max(psutil.cpu_count() - 2, 2)) -parser.add_argument('-a', '--all-faces', help='swap all faces in frame', dest='all_faces', default=False) +parser.add_argument('--max-memory', help='maximum amount of RAM in GB to be used', type=int) +parser.add_argument('--max-cores', help='number of cores to be use for CPU mode', dest='cores_count', type=int, default=max(psutil.cpu_count() - 2, 2)) +parser.add_argument('--all-faces', help='swap all faces in frame', dest='all_faces', action='store_true', default=False) for name, value in vars(parser.parse_args()).items(): args[name] = value @@ -70,8 +70,7 @@ def pre_check(): if not os.path.isfile(model_path): quit('File "inswapper_128.onnx" does not exist!') if '--gpu' in sys.argv: - NVIDIA_PROVIDERS = ['CUDAExecutionProvider', 'TensorrtExecutionProvider'] - if len(list(set(core.globals.providers) - set(NVIDIA_PROVIDERS))) == 1: + if 'ROCMExecutionProvider' not in core.globals.providers: CUDA_VERSION = torch.version.cuda CUDNN_VERSION = torch.backends.cudnn.version() if not torch.cuda.is_available() or not CUDA_VERSION: @@ -92,10 +91,6 @@ def pre_check(): def start_processing(): start_time = time.time() - threshold = len(['frame_args']) if len(args['frame_paths']) <= 10 else 10 - for i in range(threshold): - if face_check(random.choice(args['frame_paths'])) > 0.8: - quit("[WARNING] Unable to determine location of the face in the target. Please make sure the target isn't wearing clothes matching to their skin.") if args['gpu']: process_video(args['source_img'], args["frame_paths"]) end_time = time.time() @@ -200,11 +195,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 face_check(target_path) > 0.7: - quit("[WARNING] Unable to determine location of the face in the target. Please make sure the target isn't wearing clothes matching to their skin.") + if predict_image(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