added the DetectorForNSFW node.

Author: zhangp365

README.md

@@ -67,4 +67,12 @@ Added the node for convenience. The code is originally from ComfyUI-Custom-Scrip
 According to the input image ratio, decide which standard SDXL training size is the closest match. This is useful for subsequent image resizing and other processes.
 
 ## UpscaleImageWithModelIfNeed
-Enhanced the official UpscaleImageWithModel node by adding a judge. If the input image area exceeds a predefined threshold, upscaling is bypassed. The threshold is a percentage of the SDXL standard size (1024x1024) area.
+Enhanced the official UpscaleImageWithModel node by adding a judge. If the input image area exceeds a predefined threshold, upscaling is bypassed. The threshold is a percentage of the SDXL standard size (1024x1024) area.
+
+## DetectorForNSFW
+This node adapts the original model and inference code from  [nudenet](https://github.com/notAI-tech/NudeNet.git) for use with Comfy. A small 10MB default model, [320n.onnx] (https://github.com/notAI-tech/NudeNet?tab=readme-ov-file#available-models), is provided. If you wish to use other models from that repository, download the  [ONNX model] (https://github.com/notAI-tech/NudeNet?tab=readme-ov-file#available-models) and place it in the models/nsfw directory, then set the appropriate detect_size.
+
+From initial testing, the filtering effect is better than classifier models such as Falconsai/nsfw_image_detection.
+
+<img src="assets/detectorForNSFW.png" width="100%"/>
+You can also adjust the confidence levels for various rules such as buttocks_exposed to be more lenient or strict. Lower confidence levels will filter out more potential NSFW images. Setting the value to 1 will stop filtering for that specific feature.

__init__.py

@@ -3,11 +3,11 @@
 import logging
 logger = logging.getLogger(__file__)
 import os
+import importlib.util
 import shutil,filecmp
 import __main__
 
-from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS, GenderWordsConfig
-__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+from .py.nodes import GenderWordsConfig
 
 
 @server.PromptServer.instance.routes.get("/utils_node/reload_gender_words_config")
@@ -45,4 +45,34 @@ def update_javascript():
             shutil.copy(src_file, dst_file)
 
 
-update_javascript()
+update_javascript()
+
+
+NODE_CLASS_MAPPINGS = {}
+NODE_DISPLAY_NAME_MAPPINGS = {}
+
+def get_ext_dir(subpath=None, mkdir=False):
+    dir = os.path.dirname(__file__)
+    if subpath is not None:
+        dir = os.path.join(dir, subpath)
+
+    dir = os.path.abspath(dir)
+
+    if mkdir and not os.path.exists(dir):
+        os.makedirs(dir)
+    return dir
+
+py = get_ext_dir("py")
+files = os.listdir(py)
+for file in files:
+    if not file.endswith(".py"):
+        continue
+    name = os.path.splitext(file)[0]
+    imported_module = importlib.import_module(".py.{}".format(name), __name__)
+    try:
+        NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS, **imported_module.NODE_CLASS_MAPPINGS}
+        NODE_DISPLAY_NAME_MAPPINGS = {**NODE_DISPLAY_NAME_MAPPINGS, **imported_module.NODE_DISPLAY_NAME_MAPPINGS}
+    except:
+        logger.info(f"import module failure:{name}")
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

assets/detectorForNSFW.png

@@ -0,0 +1,131 @@
+import cv2
+import numpy as np
+from ..r_nudenet.nudenet import NudeDetector
+import os
+import torch
+import folder_paths as comfy_paths
+from folder_paths import models_dir
+from typing import Union, List
+import json
+import logging
+logger = logging.getLogger(__file__)
+
+comfy_paths.folder_names_and_paths["nsfw"] = ([os.path.join(models_dir, "nsfw")], {".pt",".onnx"})
+
+def tensor2np(tensor: torch.Tensor):
+    if len(tensor.shape) == 3:  # Single image
+        return np.clip(255.0 * tensor.cpu().numpy(), 0, 255).astype(np.uint8)
+    else:  # Batch of images
+        return [np.clip(255.0 * t.cpu().numpy(), 0, 255).astype(np.uint8) for t in tensor]
+
+def np2tensor(img_np: Union[np.ndarray, List[np.ndarray]]) -> torch.Tensor:
+    if isinstance(img_np, list):
+        return torch.cat([np2tensor(img) for img in img_np], dim=0)
+    return torch.from_numpy(img_np.astype(np.float32) / 255.0).unsqueeze(0)
+
+    
+class DetectorForNSFW:
+
+    def __init__(self) -> None:
+        self.model = None
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),               
+                "detect_size":([640, 320], {"default": 320}),
+                "provider": (["CPU", "CUDA", "ROCM"], ),
+            },
+            "optional": {
+                "model_name": (comfy_paths.get_filename_list("nsfw"), {"default": None}),
+                "alternative_image": ("IMAGE",),
+                "buttocks_exposed": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.05}),
+                "female_breast_exposed": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.05}),
+                "female_genitalia_exposed": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05}),
+                "anus_exposed": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.05}),
+                "male_genitalia_exposed": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE", "STRING")
+    RETURN_NAMES = ("filtered_image", "detect_result")
+    FUNCTION = "filter_exposure"
+
+    CATEGORY = "utils/filter"
+    
+    all_labels = [
+        "FEMALE_GENITALIA_COVERED",
+        "FACE_FEMALE",
+        "BUTTOCKS_EXPOSED",
+        "FEMALE_BREAST_EXPOSED",
+        "FEMALE_GENITALIA_EXPOSED",
+        "MALE_BREAST_EXPOSED",
+        "ANUS_EXPOSED",
+        "FEET_EXPOSED",
+        "BELLY_COVERED",
+        "FEET_COVERED",
+        "ARMPITS_COVERED",
+        "ARMPITS_EXPOSED",
+        "FACE_MALE",
+        "BELLY_EXPOSED",
+        "MALE_GENITALIA_EXPOSED",
+        "ANUS_COVERED",
+        "FEMALE_BREAST_COVERED",
+        "BUTTOCKS_COVERED",
+    ]
+
+    def filter_exposure(self, image, model_name=None, detect_size=320, provider="CPU", alternative_image=None, **kwargs):
+        if self.model is None:
+            self.init_model(model_name, detect_size, provider)
+
+        if alternative_image is not None:
+            alternative_image = tensor2np(alternative_image)
+        
+        images = tensor2np(image)
+        if not isinstance(images, List):
+            images = [images]
+
+        results, result_info = [],[]
+        for img in images:
+            detect_result = self.model.detect(img)
+            
+            logger.debug(f"nudenet detect result:{detect_result}")
+            filtered_results = []
+            for item in detect_result:
+                label = item['class']
+                score = item['score']
+                confidence_level = kwargs.get(label.lower())
+                if label.lower() in kwargs and score > confidence_level:
+                    filtered_results.append(item)
+            info = {"detect_result":detect_result}
+            if len(filtered_results) == 0:
+                results.append(img)
+                info["nsfw"] = False
+            else:
+                placeholder_image = alternative_image if alternative_image is not None else np.ones_like(img) * 255
+                results.append(placeholder_image)
+                info["nsfw"] = True
+
+            result_info.append(info)
+
+        result_tensor = np2tensor(results)
+        result_info = json.dumps(result_info)
+        return (result_tensor, result_info,)
+
+    def init_model(self, model_name, detect_size, provider):
+        model_path = comfy_paths.get_full_path("nsfw", model_name) if model_name else None
+        self.model = NudeDetector(model_path=model_path, providers=[provider + 'ExecutionProvider',], inference_resolution=detect_size)
+    
+
+NODE_CLASS_MAPPINGS = {
+    #image
+    "DetectorForNSFW": DetectorForNSFW,
+
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    # Image
+    "DetectorForNSFW": "detector for the NSFW",
+
+}

color_correct.py renamed to py/color_correct.py

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-utils-nodes"
-description = "Nodes:LoadImageWithSwitch, ImageBatchOneOrMore, ModifyTextGender, GenderControlOutput, ImageCompositeMaskedWithSwitch, ColorCorrectOfUtils, SplitMask, MaskFastGrow, CheckpointLoaderSimpleWithSwitch, ImageResizeTo8x, MatchImageRatioToPreset, UpscaleImageWithModelIfNeed, MaskFromFaceModel, MaskCoverFourCorners etc."
-version = "1.1.3"
+description = "Nodes:LoadImageWithSwitch, ImageBatchOneOrMore, ModifyTextGender, GenderControlOutput, ImageCompositeMaskedWithSwitch, ColorCorrectOfUtils, SplitMask, MaskFastGrow, CheckpointLoaderSimpleWithSwitch, ImageResizeTo8x, MatchImageRatioToPreset, UpscaleImageWithModelIfNeed, MaskFromFaceModel, MaskCoverFourCorners, DetectorForNSFW etc."
+version = "1.1.5"
 license = { file = "LICENSE" }
 dependencies = []