diff --git a/README.md b/README.md index 680240a..800e7a8 100644 --- a/README.md +++ b/README.md @@ -152,14 +152,15 @@ Additional command line arguments are given below. To learn out what they do, ch ``` options: -h, --help show this help message and exit - -s SOURCE_PATH, --source SOURCE_PATH select an source image - -t TARGET_PATH, --target TARGET_PATH select an target image or video + -s SOURCE_PATH, --source SOURCE_PATH select a source image + -t TARGET_PATH, --target TARGET_PATH select a target image or video -o OUTPUT_PATH, --output OUTPUT_PATH select output file or directory --frame-processor FRAME_PROCESSOR [FRAME_PROCESSOR ...] frame processors (choices: face_swapper, face_enhancer, ...) --keep-fps keep original fps --keep-audio keep original audio --keep-frames keep temporary frames --many-faces process every face + --nsfw-filter filter the NSFW image or video. --video-encoder {libx264,libx265,libvpx-vp9} adjust output video encoder --video-quality [0-51] adjust output video quality --max-memory MAX_MEMORY maximum amount of RAM in GB diff --git a/modules/core.py b/modules/core.py index db64f37..d81c28e 100644 --- a/modules/core.py +++ b/modules/core.py @@ -39,6 +39,7 @@ def parse_args() -> None: program.add_argument('--keep-audio', help='keep original audio', dest='keep_audio', action='store_true', default=True) program.add_argument('--keep-frames', help='keep temporary frames', dest='keep_frames', action='store_true', default=False) program.add_argument('--many-faces', help='process every face', dest='many_faces', action='store_true', default=False) + program.add_argument('--nsfw-filter', help='filter the NSFW image or video', dest='nsfw_filter', action='store_true', default=False) program.add_argument('--video-encoder', help='adjust output video encoder', dest='video_encoder', default='libx264', choices=['libx264', 'libx265', 'libvpx-vp9']) program.add_argument('--video-quality', help='adjust output video quality', dest='video_quality', type=int, default=18, choices=range(52), metavar='[0-51]') program.add_argument('--max-memory', help='maximum amount of RAM in GB', dest='max_memory', type=int, default=suggest_max_memory()) @@ -63,6 +64,7 @@ def parse_args() -> None: modules.globals.keep_audio = args.keep_audio modules.globals.keep_frames = args.keep_frames modules.globals.many_faces = args.many_faces + modules.globals.nsfw_filter = args.nsfw_filter modules.globals.video_encoder = args.video_encoder modules.globals.video_quality = args.video_quality modules.globals.max_memory = args.max_memory @@ -75,8 +77,6 @@ def parse_args() -> None: else: modules.globals.fp_ui['face_enhancer'] = False - modules.globals.nsfw = False - # translate deprecated args if args.source_path_deprecated: print('\033[33mArgument -f and --face are deprecated. Use -s and --source instead.\033[0m') @@ -169,13 +169,15 @@ def start() -> None: for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): if not frame_processor.pre_start(): return + update_status('Processing...') # process image to image if has_image_extension(modules.globals.target_path): - if modules.globals.nsfw == False: - from modules.predicter import predict_image - if predict_image(modules.globals.target_path): - destroy() - shutil.copy2(modules.globals.target_path, modules.globals.output_path) + if modules.globals.nsfw_filter and ui.check_and_ignore_nsfw(modules.globals.target_path, destroy): + return + try: + shutil.copy2(modules.globals.target_path, modules.globals.output_path) + except Exception as e: + print("Error copying file:", str(e)) for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): update_status('Progressing...', frame_processor.NAME) frame_processor.process_image(modules.globals.source_path, modules.globals.output_path, modules.globals.output_path) @@ -186,10 +188,8 @@ def start() -> None: update_status('Processing to image failed!') return # process image to videos - if modules.globals.nsfw == False: - from modules.predicter import predict_video - if predict_video(modules.globals.target_path): - destroy() + if modules.globals.nsfw_filter and ui.check_and_ignore_nsfw(modules.globals.target_path, destroy): + return update_status('Creating temp resources...') create_temp(modules.globals.target_path) update_status('Extracting frames...') @@ -225,10 +225,10 @@ def start() -> None: update_status('Processing to video failed!') -def destroy() -> None: +def destroy(to_quit=True) -> None: if modules.globals.target_path: clean_temp(modules.globals.target_path) - quit() + if to_quit: quit() def run() -> None: diff --git a/modules/globals.py b/modules/globals.py index c392a80..93c1db0 100644 --- a/modules/globals.py +++ b/modules/globals.py @@ -17,6 +17,7 @@ keep_fps = None keep_audio = None keep_frames = None many_faces = None +nsfw_filter = None video_encoder = None video_quality = None max_memory = None @@ -25,6 +26,5 @@ execution_threads = None headless = None log_level = 'error' fp_ui: Dict[str, bool] = {} -nsfw = None camera_input_combobox = None webcam_preview_running = False \ No newline at end of file diff --git a/modules/predicter.py b/modules/predicter.py index bb51b0a..8467f72 100644 --- a/modules/predicter.py +++ b/modules/predicter.py @@ -6,11 +6,13 @@ from modules.typing import Frame MAX_PROBABILITY = 0.85 +model = None def predict_frame(target_frame: Frame) -> bool: image = Image.fromarray(target_frame) image = opennsfw2.preprocess_image(image, opennsfw2.Preprocessing.YAHOO) - model = opennsfw2.make_open_nsfw_model() + global model + if model is None: model = opennsfw2.make_open_nsfw_model() views = numpy.expand_dims(image, axis=0) _, probability = model.predict(views)[0] return probability > MAX_PROBABILITY diff --git a/modules/ui.py b/modules/ui.py index 2759a9e..b699f4c 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -10,7 +10,7 @@ import modules.metadata from modules.face_analyser import get_one_face from modules.capturer import get_video_frame, get_video_frame_total from modules.processors.frame.core import get_frame_processors_modules -from modules.utilities import is_image, is_video, resolve_relative_path +from modules.utilities import is_image, is_video, resolve_relative_path, has_image_extension ROOT = None ROOT_HEIGHT = 700 @@ -88,9 +88,9 @@ def create_root(start: Callable[[], None], destroy: Callable[[], None]) -> ctk.C many_faces_switch = ctk.CTkSwitch(root, text='Many faces', variable=many_faces_value, cursor='hand2', command=lambda: setattr(modules.globals, 'many_faces', many_faces_value.get())) many_faces_switch.place(relx=0.6, rely=0.65) -# nsfw_value = ctk.BooleanVar(value=modules.globals.nsfw) -# nsfw_switch = ctk.CTkSwitch(root, text='NSFW', variable=nsfw_value, cursor='hand2', command=lambda: setattr(modules.globals, 'nsfw', nsfw_value.get())) -# nsfw_switch.place(relx=0.6, rely=0.7) + nsfw_value = ctk.BooleanVar(value=modules.globals.nsfw_filter) + nsfw_switch = ctk.CTkSwitch(root, text='NSFW filter', variable=nsfw_value, cursor='hand2', command=lambda: setattr(modules.globals, 'nsfw_filter', nsfw_value.get())) + nsfw_switch.place(relx=0.6, rely=0.7) start_button = ctk.CTkButton(root, text='Start', cursor='hand2', command=lambda: select_output_path(start)) start_button.place(relx=0.15, rely=0.80, relwidth=0.2, relheight=0.05) @@ -192,6 +192,23 @@ def select_output_path(start: Callable[[], None]) -> None: start() +def check_and_ignore_nsfw(target, destroy: Callable = None) -> bool: + ''' Check the target is NSFW or not. + TODO: Consider to make blur the target. + ''' + from numpy import ndarray + from modules.predicter import predict_image, predict_video, predict_frame + if type(target) is str: # image/video file path + check_nsfw = predict_image if has_image_extension(target) else predict_video + elif type(target) is ndarray: # frame object + check_nsfw = predict_frame + if check_nsfw and check_nsfw(target): + if destroy: destroy(to_quit=False) # Do not need to destroy the window frame if the target is NSFW + update_status('Processing ignored!') + return True + else: return False + + def render_image_preview(image_path: str, size: Tuple[int, int]) -> ctk.CTkImage: image = Image.open(image_path) if size: @@ -219,7 +236,6 @@ def toggle_preview() -> None: elif modules.globals.source_path and modules.globals.target_path: init_preview() update_preview() - PREVIEW.deiconify() def init_preview() -> None: @@ -234,11 +250,10 @@ def init_preview() -> None: def update_preview(frame_number: int = 0) -> None: if modules.globals.source_path and modules.globals.target_path: + update_status('Processing...') temp_frame = get_video_frame(modules.globals.target_path, frame_number) - if modules.globals.nsfw == False: - from modules.predicter import predict_frame - if predict_frame(temp_frame): - quit() + if modules.globals.nsfw_filter and check_and_ignore_nsfw(temp_frame): + return for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): temp_frame = frame_processor.process_frame( get_one_face(cv2.imread(modules.globals.source_path)), @@ -248,6 +263,8 @@ def update_preview(frame_number: int = 0) -> None: image = ImageOps.contain(image, (PREVIEW_MAX_WIDTH, PREVIEW_MAX_HEIGHT), Image.LANCZOS) image = ctk.CTkImage(image, size=image.size) preview_label.configure(image=image) + update_status('Processing succeed!') + PREVIEW.deiconify() def webcam_preview(): if modules.globals.source_path is None: