354 lines
14 KiB
Python
354 lines
14 KiB
Python
# ComfyUI-RMBG
|
|
# 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.
|
|
#
|
|
# 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": -64, "max": 64, "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)"
|
|
} |