Re-enabled the NSFW function (turn-off by default).

@refer to the PR #237 in the `experimental` branch.

Signed-off-by: Vic P <vic4key@gmail.com>

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: