From 408b0f4cf7714982945bc5a42d49a376029c32ba Mon Sep 17 00:00:00 2001 From: Vic P Date: Tue, 13 Aug 2024 04:16:34 +0700 Subject: [PATCH] ## Fix & Improve the NSFW function - Fixed incorrect state usage. - Removed the redundant argument that caused exceptions. - Prevented the app from closing when an image is flagged as NSFW. --- modules/core.py | 22 +++++++++++++++------- modules/predicter.py | 8 +++++--- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/modules/core.py b/modules/core.py index 6e2d646..1cec474 100644 --- a/modules/core.py +++ b/modules/core.py @@ -222,12 +222,18 @@ def start() -> None: def process_image_to_image() -> None: - if not modules.globals.nsfw: + if modules.globals.nsfw: from modules.predicter import predict_image if predict_image(modules.globals.target_path): - destroy() + destroy(to_quit=False) + update_status('Processing to image ignored!') + return + + try: + shutil.copy2(modules.globals.target_path, modules.globals.output_path) + except Exception as e: + print("Error copying file:", str(e)) - shutil.copy2(modules.globals.target_path, modules.globals.output_path) for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): update_status('Processing...', frame_processor.NAME) frame_processor.process_image(modules.globals.source_path, modules.globals.output_path, modules.globals.output_path) @@ -240,10 +246,12 @@ def process_image_to_image() -> None: def process_image_to_video() -> None: - if not modules.globals.nsfw: + if modules.globals.nsfw: from modules.predicter import predict_video if predict_video(modules.globals.target_path): - destroy() + destroy(to_quit=False) + update_status('Processing to video ignored!') + return update_status('Creating temporary resources...') create_temp(modules.globals.target_path) @@ -287,10 +295,10 @@ def handle_video_audio() -> None: move_temp(modules.globals.target_path, modules.globals.output_path) -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/predicter.py b/modules/predicter.py index 86c6f7d..9fd92ba 100644 --- a/modules/predicter.py +++ b/modules/predicter.py @@ -6,9 +6,11 @@ from modules.typing import Frame MAX_PROBABILITY = 0.85 # Preload the model once for efficiency -model = opennsfw2.make_open_nsfw_model() +model = None def predict_frame(target_frame: Frame) -> bool: + global model + if model is None: model = opennsfw2.make_open_nsfw_model() image = Image.fromarray(target_frame) image = opennsfw2.preprocess_image(image, opennsfw2.Preprocessing.YAHOO) views = np.expand_dims(image, axis=0) @@ -16,9 +18,9 @@ def predict_frame(target_frame: Frame) -> bool: return probability > MAX_PROBABILITY def predict_image(target_path: str) -> bool: - probability = opennsfw2.predict_image(target_path, model=model) + probability = opennsfw2.predict_image(target_path) return probability > MAX_PROBABILITY def predict_video(target_path: str) -> bool: - _, probabilities = opennsfw2.predict_video_frames(video_path=target_path, frame_interval=100, model=model) + _, probabilities = opennsfw2.predict_video_frames(video_path=target_path, frame_interval=100) return any(probability > MAX_PROBABILITY for probability in probabilities)