Merge pull request #237 from vic4key/experimental

Fix & Improve the NSFW function

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))
+
     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:

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)