nodes/1038lab-ComfyUI-RMBG
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
AI Lab
2025-05-15 14:00:14 -07:00
committed by GitHub
parent ace618ffea
commit 9a6dfcf6ea
8 changed files with 247 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files
AILab_BiRefNet.py
+22 -52
View File
@@ -331,11 +331,11 @@ class BiRefNetRMBG:
"model": "Select the BiRefNet model variant to use.",
"mask_blur": "Specify the amount of blur to apply to the mask edges (0 for no blur, higher values for more blur).",
"mask_offset": "Adjust the mask boundary (positive values expand the mask, negative values shrink it).",
"background": "Choose the background color for the final output (Alpha for transparent background).",
"invert_output": "Enable to invert both the image and mask output (useful for certain effects).",
"refine_foreground": "Use Fast Foreground Colour Estimation to optimize transparent background"
"refine_foreground": "Use Fast Foreground Colour Estimation to optimize transparent background",
"background": "Choose background type: Alpha (transparent) or Color (custom background color).",
"background_color": "Choose background color (Alpha = transparent)"
}
return {
"required": {
"image": ("IMAGE", {"tooltip": tooltips["image"]}),
@@ -344,9 +344,10 @@ class BiRefNetRMBG:
"optional": {
"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", "gray", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"refine_foreground": ("BOOLEAN", {"default": False, "tooltip": tooltips["refine_foreground"]})
"refine_foreground": ("BOOLEAN", {"default": False, "tooltip": tooltips["refine_foreground"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}),
}
}
@@ -358,35 +359,16 @@ class BiRefNetRMBG:
def process_image(self, image, model, **params):
try:
model_config = MODEL_CONFIG[model]
# Always use model's default resolution
process_res = model_config.get("default_res", 1024)
# Handle special resolution requirements
if model_config.get("force_res", False):
base_res = 512
process_res = ((process_res + base_res - 1) // base_res) * base_res
else:
process_res = process_res // 32 * 32
print(f"Using {model} model with {process_res} resolution")
params["process_res"] = process_res
processed_images = []
processed_masks = []
bg_colors = {
"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)
}
# Check and download model if needed
cache_status, message = self.model.check_model_cache(model)
if not cache_status:
print(f"Cache check: {message}")
@@ -395,18 +377,11 @@ class BiRefNetRMBG:
if not download_status:
handle_model_error(download_message)
print("Model files downloaded successfully")
# Load model if needed
self.model.load_model(model)
for img in image:
# Get mask from model
mask = self.model.process_image(img, params)
# Post-process mask
if params["mask_blur"] > 0:
mask = mask.filter(ImageFilter.GaussianBlur(radius=params["mask_blur"]))
if params["mask_offset"] != 0:
if params["mask_offset"] > 0:
for _ in range(params["mask_offset"]):
@@ -414,19 +389,12 @@ class BiRefNetRMBG:
else:
for _ in range(-params["mask_offset"]):
mask = mask.filter(ImageFilter.MinFilter(3))
if params["invert_output"]:
mask = Image.fromarray(255 - np.array(mask))
# Convert to tensors for refine_foreground
img_tensor = torch.from_numpy(np.array(tensor2pil(img))).permute(2, 0, 1).unsqueeze(0) / 255.0
mask_tensor = torch.from_numpy(np.array(mask)).unsqueeze(0).unsqueeze(0) / 255.0
if params.get("refine_foreground", False):
refined_fg = refine_foreground(
img_tensor,
mask_tensor
)
refined_fg = refine_foreground(img_tensor, mask_tensor)
refined_fg = tensor2pil(refined_fg[0].permute(1, 2, 0))
orig_image = tensor2pil(img)
r, g, b = refined_fg.split()
@@ -436,28 +404,30 @@ class BiRefNetRMBG:
orig_rgba = orig_image.convert("RGBA")
r, g, b, _ = orig_rgba.split()
foreground = Image.merge('RGBA', (r, g, b, mask))
if params["background"] != "Alpha":
bg_color = bg_colors[params["background"]]
bg_image = Image.new('RGBA', orig_image.size, (*bg_color, 255))
if params["background"] == "Alpha":
processed_images.append(pil2tensor(foreground))
else:
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
if len(hex_color) == 6:
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
a = 255
elif len(hex_color) == 8:
r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16)
else:
raise ValueError("Invalid color format")
return (r, g, b, a)
rgba = hex_to_rgba(params["background_color"])
bg_image = Image.new('RGBA', orig_image.size, rgba)
composite_image = Image.alpha_composite(bg_image, foreground)
processed_images.append(pil2tensor(composite_image.convert("RGB")))
else:
processed_images.append(pil2tensor(foreground))
processed_masks.append(pil2tensor(mask))
# Create mask image for visualization
mask_images = []
for mask_tensor in processed_masks:
# Convert mask to RGB image format for visualization
mask_image = mask_tensor.reshape((-1, 1, mask_tensor.shape[-2], mask_tensor.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
mask_images.append(mask_image)
mask_image_output = torch.cat(mask_images, dim=0)
return (torch.cat(processed_images, dim=0), torch.cat(processed_masks, dim=0), mask_image_output)
except Exception as e:
handle_model_error(f"Error in image processing: {str(e)}")
AILab_BodySegment.py
+18 -14
View File
@@ -67,8 +67,9 @@ class BodySegment:
"process_res": "Processing resolution (fixed at 512x512)",
"mask_blur": "Blur amount for mask edges",
"mask_offset": "Expand/Shrink mask boundary",
"background_color": "Choose background color (Alpha = transparent)",
"invert_output": "Invert both image and mask output",
"background": "Choose background type: Alpha (transparent) or Color (custom background color).",
"background_color": "Choose background color (Alpha = transparent)"
}
return {
@@ -80,8 +81,9 @@ class BodySegment:
for cls_name in available_classes},
"mask_blur": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1, "tooltip": tooltips["mask_blur"]}),
"mask_offset": ("INT", {"default": 0, "min": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
"background_color": (["Alpha", "black", "white", "gray", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background_color"]}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}),
},
}
@@ -121,7 +123,7 @@ class BodySegment:
except Exception as e:
return False, f"Error downloading model file: {str(e)}"
def segment_body(self, images, mask_blur=0, mask_offset=0, background_color="Alpha", invert_output=False, **class_selections):
def segment_body(self, images, mask_blur=0, mask_offset=0, background="Alpha", background_color="#222222", invert_output=False, **class_selections):
try:
# Check and download model if needed
cache_status, message = self.check_model_cache()
@@ -196,21 +198,23 @@ class BodySegment:
mask_image = Image.fromarray(255 - np.array(mask_image))
# Handle background color
if background_color == "Alpha":
if background == "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)
}
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
if len(hex_color) == 6:
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
a = 255
elif len(hex_color) == 8:
r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16)
else:
raise ValueError("Invalid color format")
return (r, g, b, a)
rgba_image = RGB2RGBA(orig_image, mask_image)
bg_image = Image.new('RGBA', orig_image.size, (*bg_colors[background_color], 255))
rgba = hex_to_rgba(background_color)
bg_image = Image.new('RGBA', orig_image.size, rgba)
composite_image = Image.alpha_composite(bg_image, rgba_image)
result_image = pil2tensor(composite_image.convert('RGB'))
AILab_ClothSegment.py
+18 -14
View File
@@ -66,8 +66,9 @@ class ClothesSegment:
"process_res": "Processing resolution (higher = more VRAM)",
"mask_blur": "Blur amount for mask edges",
"mask_offset": "Expand/Shrink mask boundary",
"background_color": "Choose background color (Alpha = transparent)",
"invert_output": "Invert both image and mask output",
"background": "Choose background type: Alpha (transparent) or Color (custom background color).",
"background_color": "Choose background color (Alpha = transparent)"
}
return {
@@ -80,8 +81,9 @@ 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": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
"background_color": (["Alpha", "black", "white", "gray", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background_color"]}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}),
},
}
@@ -141,7 +143,7 @@ class ClothesSegment:
except Exception as e:
return False, f"Error downloading model files: {str(e)}"
def segment_clothes(self, images, process_res=1024, mask_blur=0, mask_offset=0, background_color="Alpha", invert_output=False, **class_selections):
def segment_clothes(self, images, process_res=1024, mask_blur=0, mask_offset=0, background="Alpha", background_color="#222222", invert_output=False, **class_selections):
try:
# Check and download model if needed
cache_status, message = self.check_model_cache()
@@ -232,21 +234,23 @@ class ClothesSegment:
mask_image = Image.fromarray(255 - np.array(mask_image))
# Handle background color
if background_color == "Alpha":
if background == "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)
}
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
if len(hex_color) == 6:
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
a = 255
elif len(hex_color) == 8:
r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16)
else:
raise ValueError("Invalid color format")
return (r, g, b, a)
rgba_image = RGB2RGBA(orig_image, mask_image)
bg_image = Image.new('RGBA', orig_image.size, (*bg_colors[background_color], 255))
rgba = hex_to_rgba(background_color)
bg_image = Image.new('RGBA', orig_image.size, rgba)
composite_image = Image.alpha_composite(bg_image, rgba_image)
result_image = pil2tensor(composite_image.convert('RGB'))
AILab_FaceSegment.py
+18 -20
View File
@@ -59,23 +59,18 @@ class FaceSegment:
@classmethod
def INPUT_TYPES(cls):
available_classes = [
# "Background", # Not a facial feature
"Skin", "Nose", "Eyeglasses", "Left-eye", "Right-eye",
"Left-eyebrow", "Right-eyebrow", "Left-ear", "Right-ear", "Mouth",
"Upper-lip", "Lower-lip", "Hair", "Earring", "Neck",
# "Hat", # Not a facial feature
# "Necklace", # Not a facial feature
# "Clothing" # Not a facial feature
]
tooltips = {
"process_res": "Processing resolution (higher = more VRAM)",
"mask_blur": "Blur amount for mask edges",
"mask_offset": "Expand/Shrink mask boundary",
"background_color": "Choose background color (Alpha = transparent)",
"invert_output": "Invert both image and mask output",
"background": "Choose background type: Alpha (transparent) or Color (custom background color).",
"background_color": "Choose background color (Alpha = transparent)"
}
return {
"required": {
"images": ("IMAGE",),
@@ -86,8 +81,9 @@ class FaceSegment:
"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": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
"background_color": (["Alpha", "black", "white", "gray", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background_color"]}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}),
},
}
@@ -147,7 +143,7 @@ class FaceSegment:
except Exception as e:
return False, f"Error downloading model files: {str(e)}"
def segment_face(self, images, process_res=512, mask_blur=0, mask_offset=0, background_color="Alpha", invert_output=False, **class_selections):
def segment_face(self, images, process_res=512, mask_blur=0, mask_offset=0, background="Alpha", background_color="#222222", invert_output=False, **class_selections):
try:
# Check and download model if needed
cache_status, message = self.check_model_cache()
@@ -243,21 +239,23 @@ class FaceSegment:
mask_image = Image.fromarray(255 - np.array(mask_image))
# Handle background color
if background_color == "Alpha":
if background == "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)
}
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
if len(hex_color) == 6:
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
a = 255
elif len(hex_color) == 8:
r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16)
else:
raise ValueError("Invalid color format")
return (r, g, b, a)
rgba_image = RGB2RGBA(orig_image, mask_image)
bg_image = Image.new('RGBA', orig_image.size, (*bg_colors[background_color], 255))
rgba = hex_to_rgba(background_color)
bg_image = Image.new('RGBA', orig_image.size, rgba)
composite_image = Image.alpha_composite(bg_image, rgba_image)
result_image = pil2tensor(composite_image.convert('RGB'))
AILab_FashionSegment.py
+31 -18
View File
@@ -153,19 +153,30 @@ class FashionSegmentClothing:
"shoe"
]
tooltips = {
"accessories_options": "Select the accessories to be segmented",
"process_res": "Processing resolution (higher = more VRAM)",
"mask_blur": "Blur amount for mask edges",
"mask_offset": "Expand/Shrink mask boundary",
"invert_output": "Invert both image and mask output",
"background": "Choose background type: Alpha (transparent) or Color (custom background color).",
"background_color": "Choose background color (Alpha = transparent)",
}
return {
"required": {
"images": ("IMAGE",),
},
"optional": {
"accessories_options": ("ACCESSORIES_OPTIONS",),
**{cls_name: ("BOOLEAN", {"default": False})
**{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}),
"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": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}),
},
}
@@ -226,7 +237,7 @@ class FashionSegmentClothing:
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):
background="Alpha", background_color="#222222", invert_output=False, **class_selections):
try:
# Check and download model
cache_status, message = self.check_model_cache()
@@ -307,21 +318,23 @@ class FashionSegmentClothing:
mask_image = Image.fromarray(255 - np.array(mask_image))
# Process background color
if background_color == "Alpha":
if background == "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)
}
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
if len(hex_color) == 6:
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
a = 255
elif len(hex_color) == 8:
r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16)
else:
raise ValueError("Invalid color format")
return (r, g, b, a)
rgba_image = RGB2RGBA(orig_image, mask_image)
bg_image = Image.new('RGBA', orig_image.size, (*bg_colors[background_color], 255))
rgba = hex_to_rgba(background_color)
bg_image = Image.new('RGBA', orig_image.size, rgba)
composite_image = Image.alpha_composite(bg_image, rgba_image)
result_image = pil2tensor(composite_image.convert('RGB'))
AILab_RMBG.py
+19 -19
View File
@@ -1,4 +1,4 @@
# ComfyUI-RMBG v2.3.1
# ComfyUI-RMBG
# This custom node for ComfyUI provides functionality for background removal using various models,
# including RMBG-2.0, INSPYRENET, BEN, BEN2 and BIREFNET-HR. It leverages deep learning techniques
# to process images and generate masks for background removal.
@@ -545,9 +545,9 @@ class RMBG:
"process_res": "Set the processing resolution (higher values require more VRAM and may increase processing time).",
"mask_blur": "Specify the amount of blur to apply to the mask edges (0 for no blur, higher values for more blur).",
"mask_offset": "Adjust the mask boundary (positive values expand the mask, negative values shrink it).",
"background": "Choose the background color for the final output (Alpha for transparent background).",
"background": "Choose output type: Alpha (transparent) or Color (custom background color).",
"background_color": "Pick background color (supports alpha, use color picker).",
"invert_output": "Enable to invert both the image and mask output (useful for certain effects).",
"optimize": "Enable model optimization for faster processing (may affect output quality).",
"refine_foreground": "Use Fast Foreground Colour Estimation to optimize transparent background"
}
@@ -561,10 +561,10 @@ class RMBG:
"process_res": ("INT", {"default": 1024, "min": 256, "max": 2048, "step": 8, "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": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
"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"]}),
"refine_foreground": ("BOOLEAN", {"default": False, "tooltip": tooltips["refine_foreground"]})
"refine_foreground": ("BOOLEAN", {"default": False, "tooltip": tooltips["refine_foreground"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}),
}
}
@@ -578,16 +578,6 @@ class RMBG:
processed_images = []
processed_masks = []
bg_colors = {
"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)
}
model_instance = self.models[model]
# Check and download model if needed
@@ -648,9 +638,19 @@ class RMBG:
r, g, b, _ = orig_rgba.split()
foreground = Image.merge('RGBA', (r, g, b, mask))
if params["background"] != "Alpha":
bg_color = bg_colors[params["background"]]
bg_image = Image.new('RGBA', orig_image.size, (*bg_color, 255))
if params["background"] == "Color":
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
if len(hex_color) == 6:
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
a = 255
elif len(hex_color) == 8:
r, g, b, a = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16), int(hex_color[6:8], 16)
else:
raise ValueError("Invalid color format")
return (r, g, b, a)
rgba = hex_to_rgba(params["background_color"])
bg_image = Image.new('RGBA', orig_image.size, rgba)
composite_image = Image.alpha_composite(bg_image, foreground)
processed_images.append(pil2tensor(composite_image.convert("RGB")))
else:
AILab_Segment.py
+57 -29
View File
@@ -41,6 +41,18 @@ SAM_MODELS = {
"sam_vit_b (375MB)": {
"model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_vit_b.pth",
"model_type": "vit_b"
},
"sam_hq_vit_h (2.57GB)": {
"model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_hq_vit_h.pth",
"model_type": "vit_h"
},
"sam_hq_vit_l (1.25GB)": {
"model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_hq_vit_l.pth",
"model_type": "vit_l"
},
"sam_hq_vit_b (379MB)": {
"model_url": "https://huggingface.co/1038lab/sam/resolve/main/sam_hq_vit_b.pth",
"model_type": "vit_b"
}
}
@@ -114,23 +126,19 @@ def image2mask(image: Image.Image) -> torch.Tensor:
return image.squeeze()
def apply_background_color(image: Image.Image, mask_image: Image.Image,
background_color: str = "Alpha") -> Image.Image:
bg_colors = {
"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)
}
background: str = "Alpha",
background_color: str = "#222222") -> Image.Image:
rgba_image = image.copy().convert('RGBA')
rgba_image.putalpha(mask_image.convert('L'))
if background_color != "Alpha":
bg_color = bg_colors[background_color]
bg_image = Image.new('RGBA', image.size, (*bg_color, 255))
if background == "Color":
def hex_to_rgba(hex_color):
hex_color = hex_color.lstrip('#')
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
return (r, g, b, 255)
rgba = hex_to_rgba(background_color)
bg_image = Image.new('RGBA', image.size, rgba)
composite_image = Image.alpha_composite(bg_image, rgba_image)
return composite_image.convert('RGB')
@@ -144,8 +152,9 @@ class Segment:
"threshold": "Adjust mask detection strength (higher = more strict)",
"mask_blur": "Apply Gaussian blur to mask edges (0 = disabled)",
"mask_offset": "Expand/Shrink mask boundary (positive = expand, negative = shrink)",
"background_color": "Choose background color (Alpha = transparent)",
"invert_output": "Invert the mask output",
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": "Choose background type"}),
"background_color": "Choose background color (Alpha = transparent)",
}
return {
"required": {
@@ -155,16 +164,17 @@ class Segment:
"dino_model": (list(DINO_MODELS.keys()),),
},
"optional": {
"threshold": ("FLOAT", {"default": 0.35, "min": 0.05, "max": 0.95, "step": 0.01, "tooltip": tooltips["threshold"]}),
"threshold": ("FLOAT", {"default": 0.30, "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": -64, "max": 64, "step": 1, "tooltip": tooltips["mask_offset"]}),
"background_color": (["Alpha", "black", "white", "gray", "green", "blue", "red"], {"default": "Alpha", "tooltip": tooltips["background_color"]}),
"invert_output": ("BOOLEAN", {"default": False}),
"invert_output": ("BOOLEAN", {"default": False, "tooltip": tooltips["invert_output"]}),
"background": (["Alpha", "Color"], {"default": "Alpha", "tooltip": tooltips["background"]}),
"background_color": ("COLOR", {"default": "#222222", "tooltip": tooltips["background_color"]}),
}
}
RETURN_TYPES = ("IMAGE", "MASK")
RETURN_NAMES = ("IMAGE", "MASK")
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
RETURN_NAMES = ("IMAGE", "MASK", "MASK_IMAGE")
FUNCTION = "segment"
CATEGORY = "🧪AILab/🧽RMBG"
@@ -180,8 +190,8 @@ class Segment:
self.build_model = build_model
def segment(self, image, prompt, sam_model, dino_model, threshold=0.35,
mask_blur=0, mask_offset=0, background_color="Alpha",
invert_output=False):
mask_blur=0, mask_offset=0, background="Alpha",
background_color="#222222", invert_output=False):
print(f'Processing create segment for: "{prompt}"...')
image = Image.fromarray(np.clip(255. * image[0].cpu().numpy(), 0, 255).astype(np.uint8)).convert('RGBA')
@@ -193,34 +203,52 @@ class Segment:
print(f'No objects found for: "{prompt}"')
width, height = image.size
empty_mask = torch.zeros((1, height, width), dtype=torch.uint8, device="cpu")
return (empty_mask, empty_mask)
# Create empty RGB mask for visualization
empty_mask_rgb = empty_mask.reshape((-1, 1, height, width)).movedim(1, -1).expand(-1, -1, -1, 3)
return (pil2tensor(image), empty_mask, empty_mask_rgb)
masks = self.generate_masks(sam_model, image, boxes)
if masks is None:
print(f'Failed to generate mask for: "{prompt}"')
width, height = image.size
empty_mask = torch.zeros((1, height, width), dtype=torch.uint8, device="cpu")
return (empty_mask, empty_mask)
# Create empty RGB mask for visualization
empty_mask_rgb = empty_mask.reshape((-1, 1, height, width)).movedim(1, -1).expand(-1, -1, -1, 3)
return (pil2tensor(image), empty_mask, empty_mask_rgb)
mask_image = Image.fromarray((masks[1][0].numpy() * 255).astype(np.uint8))
mask_image = process_mask(mask_image, invert_output, mask_blur, mask_offset)
result_image = apply_background_color(image, mask_image, background_color)
result_image = apply_background_color(image, mask_image, background, background_color)
if background_color != "Alpha":
if background == "Color":
result_image = result_image.convert("RGB")
else:
result_image = result_image.convert("RGBA")
mask_tensor = image2mask(mask_image).unsqueeze(0)
print(f'Successfully created segment for: "{prompt}"')
return (pil2tensor(result_image), image2mask(mask_image))
# Create mask image for visualization (similar to other nodes)
mask_images = []
# Convert mask to RGB image format for visualization
mask_image_vis = mask_tensor.reshape((-1, 1, mask_image.height, mask_image.width)).movedim(1, -1).expand(-1, -1, -1, 3)
mask_images.append(mask_image_vis)
mask_image_output = torch.cat(mask_images, dim=0)
return (pil2tensor(result_image), mask_tensor, mask_image_output)
def load_sam(self, model_name):
sam_checkpoint_path = self.get_local_filepath(
SAM_MODELS[model_name]["model_url"], "sam")
model_type = SAM_MODELS[model_name]["model_type"]
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint_path)
sam = sam_model_registry[model_type]()
state_dict = torch.load(sam_checkpoint_path)
sam.load_state_dict(state_dict, strict=False)
sam_device = comfy.model_management.get_torch_device()
sam.to(device=sam_device)
sam.eval()
web/js/appearance.js
+64
View File
@@ -0,0 +1,64 @@
import { app } from "/scripts/app.js";
const COLOR_THEMES = {
segment: { nodeColor: "#222e40", nodeBgColor: "#364254", width: 320},
utility: { nodeColor: "#2e3e57", nodeBgColor: "#4b5b73"},
};
const NODE_COLORS = {
// Segmentation nodes
"RMBG": "segment",
"Segment": "segment",
"FaceSegment": "segment",
"ClothesSegment": "segment",
"BodySegment": "segment",
"FacialSegment": "segment",
"FashionSegmentAccessories": "segment",
"FashionSegmentClothing": "segment",
"BiRefNetRMBG": "segment",
// Utility nodes
"AILab_LoadImage": "utility",
"AILab_Preview": "utility",
"AILab_ImagePreview": "utility",
"AILab_MaskPreview": "utility",
"AILab_ImageMaskConvert": "utility",
"AILab_MaskEnhancer": "utility",
"AILab_MaskCombiner": "utility",
"AILab_ImageCombiner": "utility",
"AILab_MaskExtractor": "utility",
"AILab_ImageStitch": "utility",
"AILab_ImageCrop": "utility",
"AILab_ImageCropObject": "utility",
"AILab_ICLoRAConcat": "utility",
"AILab_ImageCompare": "utility",
};
function setNodeColors(node, theme) {
if (!theme) { return; }
if (theme.nodeColor) {
node.color = theme.nodeColor;
}
if (theme.nodeBgColor) {
node.bgcolor = theme.nodeBgColor;
}
if (theme.width) {
node.size = node.size || [140, 80];
node.size[0] = theme.width;
}
}
const ext = {
name: "RMBG.appearance",
nodeCreated(node) {
const nclass = node.comfyClass;
if (NODE_COLORS.hasOwnProperty(nclass)) {
let colorKey = NODE_COLORS[nclass];
const theme = COLOR_THEMES[colorKey];
setNodeColors(node, theme);
}
}
};
app.registerExtension(ext);