From 4e2be506ce72969cd799af4d55cad7c0ab4c54df Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 31 May 2023 16:30:24 +0200 Subject: [PATCH 1/8] Fix GPU for NVIDIA --- run.py | 21 ++++++++------------- 1 file changed, 8 insertions(+), 13 deletions(-) mode change 100644 => 100755 run.py diff --git a/run.py b/run.py old mode 100644 new mode 100755 index 26960dc..c74f7eb --- 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 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 = {} @@ -69,8 +69,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: @@ -89,10 +88,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() From 96a61397ef22d582fbca80cdc45fe9a782bd0071 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 31 May 2023 17:52:21 +0200 Subject: [PATCH 2/8] Implement predict video --- run.py | 56 ++++++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 52 insertions(+), 4 deletions(-) diff --git a/run.py b/run.py index c74f7eb..ae04e63 100755 --- a/run.py +++ b/run.py @@ -12,7 +12,7 @@ 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, Preprocessing from tkinter.filedialog import asksaveasfilename import webbrowser import psutil @@ -37,7 +37,7 @@ 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-memory', help='set max memory', default=16, 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)) for name, value in vars(parser.parse_args()).items(): @@ -88,6 +88,13 @@ def pre_check(): def start_processing(): start_time = time.time() + try: + seconds, probabilities = predict_video_frames(video_path=args['target_path'], frame_interval=30) + print(seconds, probabilities) + if probabilities > 0.7: + quit('0') + except: + quit('1') if args['gpu']: process_video(args['source_img'], args["frame_paths"]) end_time = time.time() @@ -140,6 +147,49 @@ def preview_video(video_path): cap.release() +def validate_video(video_path): + cap = cv2.VideoCapture('target.mp4') + frame_interval = 10 + batch_size = 10 + detector = cv2.HOGDescriptor() + + # Loop through the video frames + while True: + # Read the next batch of frames + frames = [] + for i in range(batch_size): + ret, frame = cap.read() + if not ret: + break + frames.append(frame) + + # Stop the loop if there are no more frames to process + if not frames: + break + + # Process the frames for nudity detection + for frame in frames: + # Convert the frame to grayscale + gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) + + # Apply the nudity detection algorithm + rects, weights = detector.detectMultiScale(gray, winStride=(8, 8)) + + # Draw rectangles around detected nudity regions + for (x, y, w, h) in rects: + cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 2) + + # Display the frame with detections marked + # cv2.imshow('Nudity Detection', frame) + # cv2.waitKey(1) + print('detected nude') + + # Skip ahead to the next batch of frames + cap.set(cv2.CAP_PROP_POS_MSEC, (cap.get(cv2.CAP_PROP_POS_MSEC) + frame_interval * 1000)) + + cap.release() + + def select_face(): args['source_img'] = filedialog.askopenfilename(title="Select a face") preview_image(args['source_img']) @@ -192,8 +242,6 @@ 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.") process_img(args['source_img'], target_path, args['output_file']) status("swap successful!") return From 23753b523ea34792e44afebfe1ee05c2e7b2da39 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 31 May 2023 18:00:46 +0200 Subject: [PATCH 3/8] Implement predict video --- run.py | 47 ++--------------------------------------------- 1 file changed, 2 insertions(+), 45 deletions(-) diff --git a/run.py b/run.py index ae04e63..dfcc56b 100755 --- a/run.py +++ b/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, Preprocessing +from opennsfw2 import predict_video_frames from tkinter.filedialog import asksaveasfilename import webbrowser import psutil @@ -37,7 +37,7 @@ 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', default=16, type=int) +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)) for name, value in vars(parser.parse_args()).items(): @@ -147,49 +147,6 @@ def preview_video(video_path): cap.release() -def validate_video(video_path): - cap = cv2.VideoCapture('target.mp4') - frame_interval = 10 - batch_size = 10 - detector = cv2.HOGDescriptor() - - # Loop through the video frames - while True: - # Read the next batch of frames - frames = [] - for i in range(batch_size): - ret, frame = cap.read() - if not ret: - break - frames.append(frame) - - # Stop the loop if there are no more frames to process - if not frames: - break - - # Process the frames for nudity detection - for frame in frames: - # Convert the frame to grayscale - gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) - - # Apply the nudity detection algorithm - rects, weights = detector.detectMultiScale(gray, winStride=(8, 8)) - - # Draw rectangles around detected nudity regions - for (x, y, w, h) in rects: - cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 2) - - # Display the frame with detections marked - # cv2.imshow('Nudity Detection', frame) - # cv2.waitKey(1) - print('detected nude') - - # Skip ahead to the next batch of frames - cap.set(cv2.CAP_PROP_POS_MSEC, (cap.get(cv2.CAP_PROP_POS_MSEC) + frame_interval * 1000)) - - cap.release() - - def select_face(): args['source_img'] = filedialog.askopenfilename(title="Select a face") preview_image(args['source_img']) From 54c5987ef01738aeff9858dc45feeb5478d4711c Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 31 May 2023 18:45:39 +0200 Subject: [PATCH 4/8] Implement predict video --- run.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/run.py b/run.py index dfcc56b..4e829c5 100755 --- a/run.py +++ b/run.py @@ -89,12 +89,11 @@ def pre_check(): def start_processing(): start_time = time.time() try: - seconds, probabilities = predict_video_frames(video_path=args['target_path'], frame_interval=30) - print(seconds, probabilities) - if probabilities > 0.7: - quit('0') - except: - quit('1') + seconds, probabilities = predict_video_frames(video_path=args['target_path'], frame_interval=50) + if any(probability > 0.7 for probability in probabilities): + quit() + except Exception as exception: + quit(exception) if args['gpu']: process_video(args['source_img'], args["frame_paths"]) end_time = time.time() From 4f5f79d32b653bc541ebe784c6fe9b8d981ffee5 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 31 May 2023 18:49:33 +0200 Subject: [PATCH 5/8] Implement predict video --- run.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/run.py b/run.py index 4e829c5..c530977 100755 --- a/run.py +++ b/run.py @@ -88,12 +88,9 @@ def pre_check(): def start_processing(): start_time = time.time() - try: - seconds, probabilities = predict_video_frames(video_path=args['target_path'], frame_interval=50) - if any(probability > 0.7 for probability in probabilities): - quit() - except Exception as exception: - quit(exception) + 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() From 421df6b15921223ff51f3721dd3f5806e51bcbc6 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 31 May 2023 19:06:19 +0200 Subject: [PATCH 6/8] Implement predict video --- run.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/run.py b/run.py index c530977..3a8b6b0 100755 --- a/run.py +++ b/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 From 50283c88e0f1f39648c0a6f958d2fd6a651b62ba Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 31 May 2023 19:27:31 +0200 Subject: [PATCH 7/8] Fix key error --- run.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/run.py b/run.py index 3a8b6b0..0f6bfdf 100755 --- a/run.py +++ b/run.py @@ -192,7 +192,7 @@ 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: + if predict_image(target_path) > 0.7: quit() process_img(args['source_img'], target_path, args['output_file']) status("swap successful!") From 30d33f9f5812e815bfe59c72f96374239189ed98 Mon Sep 17 00:00:00 2001 From: henryruhs Date: Wed, 31 May 2023 19:37:33 +0200 Subject: [PATCH 8/8] Improve CLI description --- README.md | 4 ++-- run.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) 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/run.py b/run.py index 0f6bfdf..3775973 100755 --- a/run.py +++ b/run.py @@ -37,8 +37,8 @@ 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('--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)) for name, value in vars(parser.parse_args()).items(): args[name] = value