From 6ec411b79a5234545701b263032d5ad3fddd98a8 Mon Sep 17 00:00:00 2001 From: Michael Wolfe Date: Tue, 29 Aug 2023 18:19:33 -0400 Subject: [PATCH] Added support for u2net, u2netp, u2net_human_seg, silueta, isnet-general-use, isnet-anime. Defaults were chosen for backwards compatability. Previous graphs should work without modification. A handful of common use colors were selected for the background colors to keep the interface cleaner. A possible future change if there is need for it. --- WAS_Node_Suite.py | 67 ++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 63 insertions(+), 4 deletions(-) diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py index 76dd8a6..a27e6e2 100644 --- a/WAS_Node_Suite.py +++ b/WAS_Node_Suite.py @@ -5651,6 +5651,10 @@ class WAS_Remove_Background: return batch # IMAGE REMBG +# Sam model needs additional input, may need to be new node entirely +# See: https://github.com/danielgatis/rembg/blob/main/USAGE.md#using-input-points +# u2net_cloth_seg model needs additional inputs, may create a new node +# An undocumented feature "putaplha" changes how alpha is applied, but does not appear to make a difference class WAS_Remove_Rembg: def __init__(self): @@ -5662,30 +5666,85 @@ class WAS_Remove_Rembg: "required": { "images": ("IMAGE",), "transparency": (["true","false"],), + "model": (["u2net", "u2netp", "u2net_human_seg", "silueta", "isnet-general-use", "isnet-anime"],), + "post_processing": ([False, True],), + "only_mask": ([False, True],), + "alpha_matting": ([False, True],), + "alpha_matting_foreground_threshold": ("INT", {"default": 240, "min": 0, "max": 255}), + "alpha_matting_background_threshold": ("INT", {"default": 10, "min": 0, "max": 255}), + "alpha_matting_erode_size": ("INT", {"default": 10, "min": 0, "max": 255}), + "background_color": (["none", "black", "white", "magenta", "chroma green", "chroma blue"],), + #"putalpha": ([False, True],), }, } + RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("images",) FUNCTION = "image_rembg" CATEGORY = "WAS Suite/Image/AI" - def image_rembg(self, images, transparency="true"): - + def image_rembg( + self, + images, + transparency="true", + model="u2net", + alpha_matting=False, + alpha_matting_foreground_threshold=240, + alpha_matting_background_threshold=10, + alpha_matting_erode_size=10, + post_processing=False, + only_mask=False, + background_color = "none", + #putalpha = False, + ): + if "rembg" not in packages(): cstr("Installing `rembg`...").msg.print() subprocess.check_call([sys.executable, '-s', '-m', 'pip', 'install', 'rembg']) - from rembg import remove + from rembg import remove, new_session os.environ['U2NET_HOME'] = os.path.join(MODELS_DIR, 'rembg') os.makedirs(os.environ['U2NET_HOME'], exist_ok=True) + # Set bgcolor + bgrgba = None + match background_color: + case "black": + bgrgba = [0, 0, 0, 255] + case "white": + bgrgba = [255, 255, 255, 255] + case "magenta": + bgrgba = [255, 0, 255, 255] + case "chroma green": + bgrgba = [0, 177, 64, 255] + case "chroma blue": + bgrgba = [0, 71, 187, 255] + case _: + bgrgba = None + + if transparency == "true" and bgrgba is not None: + bgrgba[3] = 0 + batch_tensor = [] for image in images: image = tensor2pil(image) - batch_tensor.append(pil2tensor(remove(image).convert(('RGBA' if transparency == 'true' else 'RGB')))) + batch_tensor.append(pil2tensor( + remove( + image, + session=new_session(model), + post_process_mask = post_processing, + alpha_matting = alpha_matting, + alpha_matting_foreground_threshold = alpha_matting_foreground_threshold, + alpha_matting_background_threshold = alpha_matting_background_threshold, + alpha_matting_erode_size = alpha_matting_erode_size, + only_mask = only_mask, + bgcolor = bgrgba, + #putalpha = putalpha, + ) + .convert(('RGBA' if transparency == 'true' else 'RGB')))) batch_tensor = torch.cat(batch_tensor, dim=0) return (batch_tensor,)