From c35153cbbab6c0aa1131bd49b3f033b0762b5b31 Mon Sep 17 00:00:00 2001 From: AI Lab <129358391+1038lab@users.noreply.github.com> Date: Mon, 6 Jan 2025 13:09:44 -0800 Subject: [PATCH] Add files via upload --- AILab_ClothSegment.py | 5 +- AILab_FashionSegment.py | 357 ++++++++++++++++++++++++++++++++++++++++ AILab_RMBG.py | 5 +- AILab_Segment.py | 5 +- 4 files changed, 366 insertions(+), 6 deletions(-) create mode 100644 AILab_FashionSegment.py diff --git a/AILab_ClothSegment.py b/AILab_ClothSegment.py index 13ede53..4c84807 100644 --- a/AILab_ClothSegment.py +++ b/AILab_ClothSegment.py @@ -1,4 +1,4 @@ -# ComfyUI-RMBG v1.4.0 +# ComfyUI-RMBG v1.5.0 # This custom node for ComfyUI provides functionality for background removal using various models, # including RMBG-2.0, INSPYRENET, and BEN. It leverages deep learning techniques # to process images and generate masks for background removal. @@ -83,7 +83,7 @@ class ClothesSegment: "process_res": ("INT", {"default": 512, "min": 128, "max": 2048, "step": 32, "tooltip": tooltips["process_res"]}), "mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": tooltips["mask_blur"]}), "mask_offset": ("INT", {"default": 0, "min": -20, "max": 20, "step": 1, "tooltip": tooltips["mask_offset"]}), - "background_color": (["Alpha", "black", "white", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background_color"]}), + "background_color": (["Alpha", "black", "white", "gray", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background_color"]}), "invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}), }, } @@ -242,6 +242,7 @@ class ClothesSegment: bg_colors = { "black": (0, 0, 0), "white": (255, 255, 255), + "gray": (128, 128, 128), "green": (0, 255, 0), "blue": (0, 0, 255), "red": (255, 0, 0) diff --git a/AILab_FashionSegment.py b/AILab_FashionSegment.py new file mode 100644 index 0000000..bc177f8 --- /dev/null +++ b/AILab_FashionSegment.py @@ -0,0 +1,357 @@ +# ComfyUI-RMBG v1.5.0 +# This custom node for ComfyUI provides functionality for fashion segmentation using segformer-b3-fashion model. +# It leverages deep learning techniques to process images and generate masks for fashion items segmentation. + +# Models License Notice: +# - sayeed99/segformer-b3-fashion: MIT License (https://huggingface.co/sayeed99/segformer-b3-fashion) +# +# This integration script follows GPL-3.0 License. +# When using or modifying this code, please respect both the original model licenses +# and this integration's license terms. +# +# Source: https://github.com/AILab-AI/ComfyUI-RMBG + +import os +import torch +import torch.nn as nn +import numpy as np +from typing import Tuple, Union +from PIL import Image, ImageFilter +from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation +import folder_paths +from huggingface_hub import hf_hub_download +import shutil +from torchvision import transforms + +def pil2tensor(image: Image.Image) -> torch.Tensor: + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,] + +def tensor2pil(image: torch.Tensor) -> Image.Image: + return Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8)) + +def RGB2RGBA(image: Image.Image, mask: Union[Image.Image, torch.Tensor]) -> Image.Image: + if isinstance(mask, torch.Tensor): + mask = mask2image(mask) + if mask.size != image.size: + mask = mask.resize(image.size, Image.Resampling.LANCZOS) + return Image.merge('RGBA', (*image.convert('RGB').split(), mask.convert('L'))) + +def mask2image(mask: torch.Tensor) -> Image.Image: + if len(mask.shape) == 2: + mask = mask.unsqueeze(0) + return tensor2pil(mask) + +device = "cuda" if torch.cuda.is_available() else "cpu" + +folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, "RMBG")) + +AVAILABLE_MODELS = { + "segformer_fashion": "1038lab/segformer_fashion" +} + +class FashionSegmentAccessories: + @classmethod + def INPUT_TYPES(cls): + accessories_classes = [ + # Head accessories + "hat", + "glasses", + "headband, head covering, hair accessory", + # Neck and upper body accessories + "scarf", + "tie", + # Hand accessories + "glove", + "watch", + # Waist accessories + "belt", + # Leg accessories + "leg warmer", + # Other accessories + "bag, wallet", + "umbrella" + ] + + details_classes = [ + # Upper body details + "collar", + "lapel", + "neckline", + "epaulette", + "pocket", + # Decorative details + "buckle", + "zipper", + "applique", + "bow", + "flower", + "bead", + "fringe", + "ribbon", + "rivet", + "ruffle", + "sequin", + "tassel" + ] + + return { + "required": {}, + "optional": { + **{cls_name: ("BOOLEAN", {"default": False}) + for cls_name in accessories_classes + details_classes}, + }, + } + + RETURN_TYPES = ("ACCESSORIES_OPTIONS",) + RETURN_NAMES = ("accessories_options",) + FUNCTION = "get_options" + CATEGORY = "🧪AILab/🧽RMBG" + + def get_options(self, **class_selections): + selected = [name for name, selected in class_selections.items() if selected] + return (selected,) + +class FashionSegmentClothing: + def __init__(self): + self.processor = None + self.model = None + self.cache_dir = os.path.join(folder_paths.models_dir, "RMBG", "segformer_fashion") + self.class_map = { + "Unlabelled": 0, "shirt, blouse": 1, "top, t-shirt, sweatshirt": 2, "sweater": 3, + "cardigan": 4, "jacket": 5, "vest": 6, "pants": 7, "shorts": 8, "skirt": 9, "coat": 10, + "dress": 11, "jumpsuit": 12, "cape": 13, + "glasses": 14, "hat": 15, "headband, head covering, hair accessory": 16, "tie": 17, "glove": 18, + "watch": 19, "belt": 20, "leg warmer": 21, "tights, stockings": 22, + "sock": 23, "shoe": 24, "bag, wallet": 25, "scarf": 26, "umbrella": 27, + "hood": 28, "collar": 29, "lapel": 30, "epaulette": 31, "sleeve": 32, + "pocket": 33, "neckline": 34, "buckle": 35, "zipper": 36, "applique": 37, + "bead": 38, "bow": 39, "flower": 40, "fringe": 41, "ribbon": 42, + "rivet": 43, "ruffle": 44, "sequin": 45, "tassel": 46 + } + + @classmethod + def INPUT_TYPES(cls): + clothing_classes = [ + # Upper body + "coat", + "jacket", + "cardigan", + "vest", + "sweater", + "hood", + "shirt, blouse", + "top, t-shirt, sweatshirt", + "sleeve", + # Full body + "dress", + "jumpsuit", + "cape", + # Lower body + "pants", + "shorts", + "skirt", + # Socks and shoes + "tights, stockings", + "sock", + "shoe" + ] + + return { + "required": { + "images": ("IMAGE",), + "accessories_options": ("ACCESSORIES_OPTIONS",), + }, + "optional": { + **{cls_name: ("BOOLEAN", {"default": False}) + for cls_name in clothing_classes}, + "process_res": ("INT", {"default": 512, "min": 128, "max": 2048, "step": 32}), + "mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1}), + "mask_offset": ("INT", {"default": 0, "min": -20, "max": 20, "step": 1}), + "background_color": (["Alpha", "black", "white", "gray", "green", "blue", "red"], {"default": "Alpha"}), + "invert_output": ("BOOLEAN", {"default": False}), + }, + } + + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("images", "mask") + FUNCTION = "segment_fashion" + CATEGORY = "🧪AILab/🧽RMBG" + + def check_model_cache(self): + if not os.path.exists(self.cache_dir): + return False, "Model directory not found" + + required_files = [ + 'config.json', + 'model.safetensors', + 'preprocessor_config.json' + ] + + missing_files = [f for f in required_files if not os.path.exists(os.path.join(self.cache_dir, f))] + if missing_files: + return False, f"Missing required model files: {', '.join(missing_files)}" + return True, "Model cache verified" + + def clear_model(self): + if self.model is not None: + self.model.cpu() + del self.model + self.model = None + self.processor = None + torch.cuda.empty_cache() + + def download_model_files(self): + model_id = AVAILABLE_MODELS["segformer_fashion"] + model_files = { + 'config.json': 'config.json', + 'model.safetensors': 'model.safetensors', + 'preprocessor_config.json': 'preprocessor_config.json' + } + + os.makedirs(self.cache_dir, exist_ok=True) + print(f"Downloading fashion segmentation model files...") + + try: + for save_name, repo_path in model_files.items(): + print(f"Downloading {save_name}...") + downloaded_path = hf_hub_download( + repo_id=model_id, + filename=repo_path, + local_dir=self.cache_dir, + local_dir_use_symlinks=False + ) + + if os.path.dirname(downloaded_path) != self.cache_dir: + target_path = os.path.join(self.cache_dir, save_name) + shutil.move(downloaded_path, target_path) + return True, "Model files downloaded successfully" + except Exception as e: + return False, f"Error downloading model files: {str(e)}" + + def segment_fashion(self, images, accessories_options, process_res=512, mask_blur=0, mask_offset=0, + background_color="Alpha", invert_output=False, **class_selections): + try: + # Check and download model + cache_status, message = self.check_model_cache() + if not cache_status: + print(f"Cache check: {message}") + download_status, download_message = self.download_model_files() + if not download_status: + raise RuntimeError(download_message) + + # Load model if needed + if self.processor is None: + self.processor = SegformerImageProcessor.from_pretrained(self.cache_dir) + self.model = AutoModelForSemanticSegmentation.from_pretrained(self.cache_dir) + self.model.eval() + for param in self.model.parameters(): + param.requires_grad = False + self.model.to(device) + + # Get all selected classes + selected_classes = [] + # Add clothing selections + selected_classes.extend([name for name, selected in class_selections.items() if selected]) + # Add accessories selections + selected_classes.extend(accessories_options) + + if not selected_classes: + selected_classes = ["shirt, blouse"] + + transform_image = transforms.Compose([ + transforms.Resize((process_res, process_res)), + transforms.ToTensor(), + ]) + + batch_tensor = [] + batch_masks = [] + + for image in images: + orig_image = tensor2pil(image) + w, h = orig_image.size + + inputs = self.processor(images=orig_image, return_tensors="pt") + inputs = {k: v.to(device) for k, v in inputs.items()} + + with torch.no_grad(): + outputs = self.model(**inputs) + logits = outputs.logits.cpu() + + upsampled_logits = nn.functional.interpolate( + logits, + size=(h, w), + mode="bilinear", + align_corners=False, + ) + pred_seg = upsampled_logits.argmax(dim=1)[0] + + # Merge masks for selected classes + combined_mask = None + for class_name in selected_classes: + mask = (pred_seg == self.class_map[class_name]).float() + if combined_mask is None: + combined_mask = mask + else: + combined_mask = torch.clamp(combined_mask + mask, 0, 1) + + # Convert mask to PIL for processing + mask_image = Image.fromarray((combined_mask.numpy() * 255).astype(np.uint8)) + + if mask_blur > 0: + mask_image = mask_image.filter(ImageFilter.GaussianBlur(radius=mask_blur)) + + if mask_offset != 0: + if mask_offset > 0: + mask_image = mask_image.filter(ImageFilter.MaxFilter(size=mask_offset * 2 + 1)) + else: + mask_image = mask_image.filter(ImageFilter.MinFilter(size=-mask_offset * 2 + 1)) + + if invert_output: + mask_image = Image.fromarray(255 - np.array(mask_image)) + + # Process background color + if background_color == "Alpha": + rgba_image = RGB2RGBA(orig_image, mask_image) + result_image = pil2tensor(rgba_image) + else: + bg_colors = { + "black": (0, 0, 0), + "white": (255, 255, 255), + "gray": (128, 128, 128), + "green": (0, 255, 0), + "blue": (0, 0, 255), + "red": (255, 0, 0) + } + + rgba_image = RGB2RGBA(orig_image, mask_image) + bg_image = Image.new('RGBA', orig_image.size, (*bg_colors[background_color], 255)) + composite_image = Image.alpha_composite(bg_image, rgba_image) + result_image = pil2tensor(composite_image.convert('RGB')) + + batch_tensor.append(result_image) + batch_masks.append(pil2tensor(mask_image)) + + batch_tensor = torch.cat(batch_tensor, dim=0) + batch_masks = torch.cat(batch_masks, dim=0) + + return (batch_tensor, batch_masks) + + except Exception as e: + self.clear_model() + raise RuntimeError(f"Error in fashion segmentation: {str(e)}") + finally: + if not self.model.training: + self.clear_model() + + def __del__(self): + self.clear_model() + +NODE_CLASS_MAPPINGS = { + "FashionSegmentAccessories": FashionSegmentAccessories, + "FashionSegmentClothing": FashionSegmentClothing +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "FashionSegmentAccessories": "Accessories Segment (RMBG)", + "FashionSegmentClothing": "Fashion Segment (RMBG)" +} \ No newline at end of file diff --git a/AILab_RMBG.py b/AILab_RMBG.py index d868c00..b862783 100644 --- a/AILab_RMBG.py +++ b/AILab_RMBG.py @@ -1,4 +1,4 @@ -# ComfyUI-RMBG v1.4.0 +# ComfyUI-RMBG v1.5.0 # This custom node for ComfyUI provides functionality for background removal using various models, # including RMBG-2.0, INSPYRENET, and BEN. It leverages deep learning techniques # to process images and generate masks for background removal. @@ -335,7 +335,7 @@ class RMBG: "process_res": ("INT", {"default": 1024, "min": 256, "max": 2048, "step": 32, "tooltip": tooltips["process_res"]}), "mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": tooltips["mask_blur"]}), "mask_offset": ("INT", {"default": 0, "min": -20, "max": 20, "step": 1, "tooltip": tooltips["mask_offset"]}), - "background": (["Alpha", "black", "white", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background"]}), + "background": (["Alpha", "black", "white", "gray", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background"]}), "invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}), "optimize": (["default", "on"], {"default": "default", "tooltip": tooltips["optimize"]}) } @@ -355,6 +355,7 @@ class RMBG: "Alpha": None, "black": (0, 0, 0), "white": (255, 255, 255), + "gray": (128, 128, 128), "green": (0, 255, 0), "blue": (0, 0, 255), "red": (255, 0, 0) diff --git a/AILab_Segment.py b/AILab_Segment.py index f006332..d70f372 100644 --- a/AILab_Segment.py +++ b/AILab_Segment.py @@ -1,4 +1,4 @@ -# ComfyUI-RMBG v1.4.0 +# ComfyUI-RMBG v1.5.0 # This custom node for ComfyUI provides functionality for background removal using various models, # including RMBG-2.0, INSPYRENET, and BEN. It leverages deep learning techniques # to process images and generate masks for background removal. @@ -119,6 +119,7 @@ def apply_background_color(image: Image.Image, mask_image: Image.Image, "Alpha": None, "black": (0, 0, 0), "white": (255, 255, 255), + "gray": (128, 128, 128), "green": (0, 255, 0), "blue": (0, 0, 255), "red": (255, 0, 0) @@ -157,7 +158,7 @@ class Segment: "threshold": ("FLOAT", {"default": 0.35, "min": 0.05, "max": 0.95, "step": 0.01, "tooltip": tooltips["threshold"]}), "mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": tooltips["mask_blur"]}), "mask_offset": ("INT", {"default": 0, "min": -20, "max": 20, "step": 1, "tooltip": tooltips["mask_offset"]}), - "background_color": (["Alpha", "black", "white", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background_color"]}), + "background_color": (["Alpha", "black", "white", "gray", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background_color"]}), "invert_output": ("BOOLEAN", {"default": False}), } }