From 9a6dfcf6eaf8d816ab2e444b75bc5a33f93f9a8c Mon Sep 17 00:00:00 2001 From: AI Lab <129358391+1038lab@users.noreply.github.com> Date: Thu, 15 May 2025 14:00:14 -0700 Subject: [PATCH] Add files via upload --- AILab_BiRefNet.py | 74 +++++++++++------------------------ AILab_BodySegment.py | 32 ++++++++------- AILab_ClothSegment.py | 32 ++++++++------- AILab_FaceSegment.py | 38 +++++++++--------- AILab_FashionSegment.py | 49 ++++++++++++++--------- AILab_RMBG.py | 38 +++++++++--------- AILab_Segment.py | 86 +++++++++++++++++++++++++++-------------- web/js/appearance.js | 64 ++++++++++++++++++++++++++++++ 8 files changed, 247 insertions(+), 166 deletions(-) create mode 100644 web/js/appearance.js diff --git a/AILab_BiRefNet.py b/AILab_BiRefNet.py index 6465fbe..f1320f0 100644 --- a/AILab_BiRefNet.py +++ b/AILab_BiRefNet.py @@ -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)}") diff --git a/AILab_BodySegment.py b/AILab_BodySegment.py index 29dedb1..9b2ac84 100644 --- a/AILab_BodySegment.py +++ b/AILab_BodySegment.py @@ -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')) diff --git a/AILab_ClothSegment.py b/AILab_ClothSegment.py index d311f37..742431b 100644 --- a/AILab_ClothSegment.py +++ b/AILab_ClothSegment.py @@ -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')) diff --git a/AILab_FaceSegment.py b/AILab_FaceSegment.py index acd15d5..3d53a7b 100644 --- a/AILab_FaceSegment.py +++ b/AILab_FaceSegment.py @@ -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')) diff --git a/AILab_FashionSegment.py b/AILab_FashionSegment.py index d37f524..3a2f264 100644 --- a/AILab_FashionSegment.py +++ b/AILab_FashionSegment.py @@ -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')) diff --git a/AILab_RMBG.py b/AILab_RMBG.py index 27797fb..0858526 100644 --- a/AILab_RMBG.py +++ b/AILab_RMBG.py @@ -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: diff --git a/AILab_Segment.py b/AILab_Segment.py index 92b6008..c10b543 100644 --- a/AILab_Segment.py +++ b/AILab_Segment.py @@ -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() diff --git a/web/js/appearance.js b/web/js/appearance.js new file mode 100644 index 0000000..7b73929 --- /dev/null +++ b/web/js/appearance.js @@ -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);