import torch import numpy as np from PIL import Image, ImageFilter import cv2 import os import folder_paths import logging import urllib.request try: from ultralytics import YOLO except ImportError: logging.error("Ultralytics is not installed. Please run `pip install ultralytics`.") # Models directory for YOLO/Ultralytics ultralytics_dir = os.path.join(folder_paths.models_dir, "ultralytics") yolo_dir = os.path.join(folder_paths.models_dir, "yolo") # If ultralytics directory exists, prefer it, otherwise use yolo (create it if necessary) target_model_dir = ultralytics_dir if os.path.exists(ultralytics_dir) else yolo_dir if not os.path.exists(target_model_dir): try: os.makedirs(target_model_dir) except Exception as e: logging.error(f"Failed to create models directory {target_model_dir}: {e}") CUSTOM_MODEL_URL = "https://huggingface.co/sugarknight/sensitive-detect/resolve/main/sensitive_detect_v07.pt" CUSTOM_MODEL_NAME = "sensitive_detect_v07.pt" NO_SECOND_MODEL = "None" # YOLO model file extensions MODEL_EXTENSIONS = (".pt", ".pth", ".onnx") def get_available_models(): """Scan model directories and return a list of available model filenames.""" available = set() # Always include the custom model as an option available.add(CUSTOM_MODEL_NAME) # Scan both directories for model files for d in [ultralytics_dir, yolo_dir]: if os.path.isdir(d): for f in os.listdir(d): if f.lower().endswith(MODEL_EXTENSIONS) and os.path.isfile(os.path.join(d, f)): available.add(f) model_list = sorted(available) # Ensure custom model is first (default) if CUSTOM_MODEL_NAME in model_list: model_list.remove(CUSTOM_MODEL_NAME) model_list.insert(0, CUSTOM_MODEL_NAME) return model_list if model_list else [CUSTOM_MODEL_NAME] def download_file(url, target_path): print(f"Downloading {url} to {target_path}...") try: # User-Agent header is important for some CDNs req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'}) with urllib.request.urlopen(req) as response: with open(target_path, 'wb') as out_file: # Add simple block reading total_length = response.getheader('content-length') if total_length is None: # no content length header out_file.write(response.read()) else: total_length = int(total_length) downloaded = 0 block_size = max(4096, total_length // 100) while True: buffer = response.read(block_size) if not buffer: break downloaded += len(buffer) out_file.write(buffer) print(f"Successfully downloaded {CUSTOM_MODEL_NAME}.") return True except Exception as e: print(f"Error downloading {CUSTOM_MODEL_NAME}: {e}") return False try: from .extract_segments import ( extract_and_save_segments, MosaicMode, detect_frame, detect_frame_with_models, render_processed_layer, compose_frame, fill_video_gaps, save_images_to_psd, ) except ImportError: from extract_segments import ( extract_and_save_segments, MosaicMode, detect_frame, detect_frame_with_models, render_processed_layer, compose_frame, fill_video_gaps, save_images_to_psd, ) class AutoMosaic: def __init__(self): self.models = {} @classmethod def INPUT_TYPES(s): available_models = get_available_models() second_model_choices = [NO_SECOND_MODEL] + available_models return { "required": { "image": ("IMAGE",), "save_psd": ("BOOLEAN", {"default": False}), "filename_prefix": ("STRING", {"default": "AutoMosaic"}), "confidence": ("FLOAT", {"default": 0.25, "min": 0.01, "max": 1.0, "step": 0.01}), "process_method": (["raw", "mosaic", "white", "blur"], {"default": "mosaic"}), "factor": ("INT", {"default": 100, "min": 10, "step": 1}), "target_class": ("STRING", {"default": "pussy,penis"}), "model_name": (available_models, {"default": CUSTOM_MODEL_NAME}), "mask_expand": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 20.0, "step": 0.1}), "video_mode": ("BOOLEAN", {"default": False}), "buffer_frames": ("INT", {"default": 10, "min": 1, "max": 120, "step": 1}), "morph_method": (["simple", "optical_flow"], {"default": "simple"}), "model_name_2": (second_model_choices, {"default": NO_SECOND_MODEL}), }, } RETURN_TYPES = ("IMAGE", "MASK") OUTPUT_NODE = True FUNCTION = "process_image" CATEGORY = "image/process" def _ensure_model(self, model_name): if model_name in self.models: return self.models[model_name] potential_paths = [ os.path.join(ultralytics_dir, model_name), os.path.join(yolo_dir, model_name), ] model_path = model_name for path in potential_paths: if os.path.isfile(path): model_path = path break if not os.path.isfile(model_path) and model_name == CUSTOM_MODEL_NAME: download_dest = os.path.join(target_model_dir, CUSTOM_MODEL_NAME) if download_file(CUSTOM_MODEL_URL, download_dest): model_path = download_dest logging.info(f"Loading YOLO model: {model_path}") model = YOLO(model_path) self._warmup_model(model, model_name) self.models[model_name] = model return model @staticmethod def _warmup_model(model, model_name): logging.info(f"Warming up YOLO model: {model_name}") warmup_image = Image.new("RGB", (64, 64), (0, 0, 0)) model.predict(warmup_image, conf=0.99, verbose=False) def _ensure_models(self, model_name, model_name_2=NO_SECOND_MODEL): model_names = [] for name in (model_name, model_name_2): if name and name != NO_SECOND_MODEL and name not in model_names: model_names.append(name) if not model_names: model_names.append(CUSTOM_MODEL_NAME) return [(name, self._ensure_model(name)) for name in model_names] @staticmethod def _resolve_mosaic_mode(process_method): if process_method == "mosaic": return MosaicMode.MOSAIC if process_method == "blur": return MosaicMode.BLUR if process_method == "white": return MosaicMode.WHITE return MosaicMode.RAW def process_image(self, image, save_psd, filename_prefix, confidence, process_method, factor, target_class, model_name=CUSTOM_MODEL_NAME, mask_expand=0.0, video_mode=False, buffer_frames=10, morph_method="simple", model_name_2=NO_SECOND_MODEL): active_models = self._ensure_models(model_name, model_name_2) output_dir = folder_paths.get_output_directory() mosaic_mode = self._resolve_mosaic_mode(process_method) if video_mode and len(image) > 1: return self._process_video( image=image, save_psd=save_psd, filename_prefix=filename_prefix, confidence=confidence, mosaic_mode=mosaic_mode, factor=factor, target_class=target_class, output_dir=output_dir, mask_expand=mask_expand, buffer_frames=buffer_frames, morph_method=morph_method, active_models=active_models, ) processed_images = [] processed_masks = [] for img_tensor in image: img_np = (img_tensor.cpu().numpy() * 255).astype(np.uint8) img_pil = Image.fromarray(img_np) full_output_folder, filename, counter, _subfolder, _prefix_res = folder_paths.get_save_image_path( filename_prefix, output_dir, img_np.shape[1], img_np.shape[0] ) base_name = f"{filename}_{counter:05}" composite_pil, mask_pil = extract_and_save_segments( img_pil=img_pil, base_name=base_name, model=active_models[0][1], output_dir=full_output_folder, confidence=confidence, mosaic_mode=mosaic_mode, target_class=target_class, block_size=factor, save_psd=save_psd, mask_expand=mask_expand, models=active_models, ) out_img_np = np.array(composite_pil).astype(np.float32) / 255.0 processed_images.append(torch.from_numpy(out_img_np)) mask_np = np.array(mask_pil).astype(np.float32) / 255.0 processed_masks.append(torch.from_numpy(mask_np)) return {"ui": {"text": ["Successfully saved PSD"]}, "result": (torch.stack(processed_images), torch.stack(processed_masks))} def _process_video(self, image, save_psd, filename_prefix, confidence, mosaic_mode, factor, target_class, output_dir, mask_expand, buffer_frames, morph_method, active_models): target_classes = [c.strip() for c in target_class.split(",")] frames_pil = [] frames_np = [] for img_tensor in image: arr = (img_tensor.cpu().numpy() * 255).astype(np.uint8) frames_np.append(arr) frames_pil.append(Image.fromarray(arr)) h, w = frames_np[0].shape[:2] full_output_folder, filename, counter, _subfolder, _prefix_res = folder_paths.get_save_image_path( filename_prefix, output_dir, w, h ) os.makedirs(full_output_folder, exist_ok=True) # Detect on every frame. frame_class_masks = [] for img_pil in frames_pil: class_masks, _segments = detect_frame_with_models( img_pil, active_models, confidence, target_classes, factor, mask_expand=mask_expand, ) frame_class_masks.append(class_masks) # Fill gaps using simple morph + optical-flow tracking per spec. fill_video_gaps(frames_np, frame_class_masks, buffer_frames, morph_method) # Render each frame. processed_images = [] processed_masks = [] for i, img_pil in enumerate(frames_pil): layers_rgba = [] combined = np.zeros((h, w), dtype=np.float32) for cls_name, mask in frame_class_masks[i].items(): combined = np.maximum(combined, mask) layers_rgba.append(render_processed_layer(img_pil, mask, mosaic_mode, factor)) composite_pil = compose_frame(img_pil, layers_rgba) if save_psd: psd_layers = [{"name": "base image", "image": img_pil}] for idx, (cls_name, mask) in enumerate(frame_class_masks[i].items()): psd_layers.append({"name": f"{cls_name} image {idx}", "image": layers_rgba[idx]}) base_name = f"{filename}_{counter + i:05}" add_index = 0 save_path = os.path.join(full_output_folder, f"{base_name}.psd") while os.path.exists(save_path): add_index += 1 save_path = os.path.join(full_output_folder, f"{base_name}_{add_index:05}.psd") save_images_to_psd(psd_layers, save_path) out_img_np = np.array(composite_pil).astype(np.float32) / 255.0 processed_images.append(torch.from_numpy(out_img_np)) processed_masks.append(torch.from_numpy(np.clip(combined, 0, 1).astype(np.float32))) return {"ui": {"text": [f"Processed {len(frames_pil)} frames in video mode"]}, "result": (torch.stack(processed_images), torch.stack(processed_masks))} # Standard ComfyUI node mappings NODE_CLASS_MAPPINGS = { "AutoMosaic": AutoMosaic } NODE_DISPLAY_NAME_MAPPINGS = { "AutoMosaic": "Auto Mosaic" }