diff --git a/README.md b/README.md index f252fb8..09283ab 100644 --- a/README.md +++ b/README.md @@ -40,7 +40,7 @@ The following models are supported: ## Nodes -### BiRefNetUltraV3_RBS +### BiRefNetUltra_RBS - **Input**: Image - **Output**: Mask - **Parameters**: @@ -65,6 +65,13 @@ The following models are supported: ## Changelog +### v1.2.0 +- 为 BiRefNet Ultra (RBS) 和 Transparent Background Ultra (RBS) 节点添加了 mask 输入功能 +- 新增 "自动适用原图遮罩" 参数，支持智能检测和使用输入的 mask +- 优化了节点名称，移除了 BiRefNetUltraV3_RBS 中的 "V3" 后缀 +- 修复了输入 mask 应用时的翻转问题，确保输出结果正确 +- 改进了批量处理逻辑，支持图像和 mask 的自动匹配 + ### v1.1.2 - 修复了模型加载路径的问题 - 改进了模型自动下载功能 diff --git a/__init__.py b/__init__.py index 504a6f6..6e18409 100644 --- a/__init__.py +++ b/__init__.py @@ -2,4 +2,4 @@ from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] -__version__ = 'v1.1.1' \ No newline at end of file +__version__ = 'v1.2.0' \ No newline at end of file diff --git a/nodes.py b/nodes.py index 039ea97..93ec52a 100644 --- a/nodes.py +++ b/nodes.py @@ -1,6 +1,6 @@ import torch import numpy as np -from PIL import Image +from PIL import Image, ImageEnhance import folder_paths import os from .model_utils import load_model, preprocess_image, postprocess_mask, apply_transparency @@ -45,8 +45,10 @@ class TransparentBackgroundUltra_RBS: "model": (list(scan_transparent_models().keys()),), "device": (device_list,), "max_megapixels": ("FLOAT", {"default": 2.0, "min": 1, "max": 999, "step": 0.1}), + "auto_use_original_mask": ("BOOLEAN", {"default": True}), }, "optional": { + "mask": ("MASK",), } } @@ -54,55 +56,107 @@ class TransparentBackgroundUltra_RBS: RETURN_NAMES = ("image", "mask", ) FUNCTION = "transparent_background_ultra" CATEGORY = 'RemoveBackgroundSuite' + + def is_mask_valid(self, mask): + """检查mask是否有效（非空且有内容）""" + if mask is None: + return False + + # 检查mask是否为tensor + if isinstance(mask, torch.Tensor): + # 检查是否有非零值 + if mask.numel() == 0: + return False + # 检查是否有非零像素 + if torch.sum(mask > 0.01).item() == 0: # 使用0.01作为阈值避免浮点误差 + return False + return True + + return False - def transparent_background_ultra(self, image, model, device, max_megapixels): + def transparent_background_ultra(self, image, model, device, max_megapixels, auto_use_original_mask, mask=None): import glob from transparent_background import Remover mode_dict = {"ckpt_base.pth": "base", "ckpt_base_nightly.pth": "base-nightly", "ckpt_fast.pth": "fast"} ret_images = [] ret_masks = [] - model_file_list = glob.glob(os.path.join(folder_paths.models_dir, "transparent-background") + '/*.pth') - model_dict = {} - for i in range(len(model_file_list)): - _, __filename = os.path.split(model_file_list[i]) - model_dict[__filename] = model_file_list[i] - try: - mode = mode_dict[model] - except: - mode = "base" - remover = Remover(mode=mode, jit=False, device=device, ckpt=model_dict[model]) - for i in image: - i = torch.unsqueeze(i, 0) - orig_image = tensor2pil(i).convert('RGB') + + # 检查是否使用输入的mask + use_input_mask = False + if auto_use_original_mask and self.is_mask_valid(mask): + use_input_mask = True + log(f"[TransparentBackgroundUltra_RBS] 检测到有效输入mask，将使用输入mask进行处理", message_type='info') + else: + # 正常加载模型 + model_file_list = glob.glob(os.path.join(folder_paths.models_dir, "transparent-background") + '/*.pth') + model_dict = {} + for i in range(len(model_file_list)): + _, __filename = os.path.split(model_file_list[i]) + model_dict[__filename] = model_file_list[i] + try: + mode = mode_dict[model] + except: + mode = "base" + remover = Remover(mode=mode, jit=False, device=device, ckpt=model_dict[model]) + + for i, img in enumerate(image): + img = torch.unsqueeze(img, 0) + orig_image = tensor2pil(img).convert('RGB') width, height = orig_image.size - max_pixels = int(max_megapixels * 1_048_576) - orig_pixels = width * height - patch_size = 32 - if orig_pixels > max_pixels: - scale = (max_pixels / orig_pixels) ** 0.5 - new_width = max(1, int(width * scale)) - new_height = max(1, int(height * scale)) + + if use_input_mask: + # 使用输入的mask + if i < mask.shape[0]: # 确保有对应的mask + input_mask = mask[i:i+1] # 取对应的mask + else: + # 如果mask数量不足，使用最后一个mask + input_mask = mask[-1:] + + # 将mask转换为PIL图像并调整大小 + mask_pil = tensor2pil(input_mask) + if mask_pil.mode != 'L': + mask_pil = mask_pil.convert('L') + if mask_pil.size != orig_image.size: + mask_pil = mask_pil.resize(orig_image.size, Image.BILINEAR) + + # 翻转mask以确保正确的结果 + mask_pil = ImageOps.invert(mask_pil) + + # 直接使用输入mask创建透明图像 + ret_image = RGB2RGBA(orig_image, mask_pil) + _mask = mask_pil + log(f"[TransparentBackgroundUltra_RBS] 使用输入mask，跳过模型推理", message_type='info') else: - new_width, new_height = width, height - new_width = (new_width // patch_size) * patch_size - new_height = (new_height // patch_size) * patch_size - new_width = max(patch_size, new_width) - new_height = max(patch_size, new_height) - inference_image_size = (new_width, new_height) - log(f"[TransparentBackgroundUltra_RBS] 原始尺寸: {width}x{height}, 实际推理尺寸: {inference_image_size[0]}x{inference_image_size[1]}", message_type='info') - resized_image = orig_image.resize(inference_image_size, Image.BILINEAR) - ret_image = remover.process(resized_image, type='rgba') - # 推理后还原为原图尺寸 - ret_image = ret_image.resize(orig_image.size, Image.BILINEAR) - _mask = ret_image.split()[3] + # 正常进行模型推理 + max_pixels = int(max_megapixels * 1_048_576) + orig_pixels = width * height + patch_size = 32 + if orig_pixels > max_pixels: + scale = (max_pixels / orig_pixels) ** 0.5 + new_width = max(1, int(width * scale)) + new_height = max(1, int(height * scale)) + else: + new_width, new_height = width, height + new_width = (new_width // patch_size) * patch_size + new_height = (new_height // patch_size) * patch_size + new_width = max(patch_size, new_width) + new_height = max(patch_size, new_height) + inference_image_size = (new_width, new_height) + log(f"[TransparentBackgroundUltra_RBS] 原始尺寸: {width}x{height}, 实际推理尺寸: {inference_image_size[0]}x{inference_image_size[1]}", message_type='info') + resized_image = orig_image.resize(inference_image_size, Image.BILINEAR) + ret_image = remover.process(resized_image, type='rgba') + # 推理后还原为原图尺寸 + ret_image = ret_image.resize(orig_image.size, Image.BILINEAR) + _mask = ret_image.split()[3] + ret_images.append(pil2tensor(ret_image)) ret_masks.append(image2mask(_mask)) log(f"{self.NODE_NAME} Processed {len(ret_images)} image(s).", message_type='finish') return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),) -# ======================== V3 新增：支持 BiRefNet-Dynamic ======================== +# ======================== 支持 BiRefNet-Dynamic ======================== -class BiRefNetUltraV3_RBS: +class BiRefNetUltra_RBS: birefnet_model_repos = { "BiRefNet-General": "ZhengPeng7/BiRefNet", "RMBG-2.0": "briaai/RMBG-2.0", @@ -112,7 +166,7 @@ class BiRefNetUltraV3_RBS: } def __init__(self): - self.NODE_NAME = 'BiRefNetUltraV3_RBS' + self.NODE_NAME = 'BiRefNetUltra_RBS' self.model_cache = {} @classmethod @@ -125,14 +179,16 @@ class BiRefNetUltraV3_RBS: "version": (model_list, {"default": model_list[0]}), "device": (device_list,), "max_megapixels": ("FLOAT", {"default": 2.0, "min": 1, "max": 999, "step": 0.1}), + "auto_use_original_mask": ("BOOLEAN", {"default": True}), }, "optional": { + "mask": ("MASK",), } } RETURN_TYPES = ("IMAGE", "MASK", ) RETURN_NAMES = ("image", "mask", ) - FUNCTION = "birefnet_ultra_v3" + FUNCTION = "birefnet_ultra" CATEGORY = 'RemoveBackgroundSuite' def load_birefnet_model(self, version): @@ -169,66 +225,119 @@ class BiRefNetUltraV3_RBS: else: raise RuntimeError(f"Model path {model_path} does not exist") - def birefnet_ultra_v3(self, image, version, device, max_megapixels): + def is_mask_valid(self, mask): + """检查mask是否有效（非空且有内容）""" + if mask is None: + return False + + # 检查mask是否为tensor + if isinstance(mask, torch.Tensor): + # 检查是否有非零值 + if mask.numel() == 0: + return False + # 检查是否有非零像素 + if torch.sum(mask > 0.01).item() == 0: # 使用0.01作为阈值避免浮点误差 + return False + return True + + return False + + def birefnet_ultra(self, image, version, device, max_megapixels, auto_use_original_mask, mask=None): ret_images = [] ret_masks = [] torch.set_float32_matmul_precision(['high', 'highest'][0]) - # 模型缓存，避免重复加载 - if version not in self.model_cache: - self.model_cache[version] = self.load_birefnet_model(version) - birefnet_model = self.model_cache[version] - birefnet_model.to(device) - birefnet_model.eval() + + # 检查是否使用输入的mask + use_input_mask = False + if auto_use_original_mask and self.is_mask_valid(mask): + use_input_mask = True + log(f"[BiRefNetUltra_RBS] 检测到有效输入mask，将使用输入mask进行处理", message_type='info') + else: + # 模型缓存，避免重复加载 + if version not in self.model_cache: + self.model_cache[version] = self.load_birefnet_model(version) + birefnet_model = self.model_cache[version] + birefnet_model.to(device) + birefnet_model.eval() + comfy_pbar = ProgressBar(len(image)) - tqdm_pbar = tqdm.tqdm(total=len(image), desc="Processing BiRefNetV3") - for i in image: - i = torch.unsqueeze(i, 0) - orig_image = tensor2pil(i).convert('RGB') + tqdm_pbar = tqdm.tqdm(total=len(image), desc="Processing BiRefNet") + + for i, img in enumerate(image): + img = torch.unsqueeze(img, 0) + orig_image = tensor2pil(img).convert('RGB') width, height = orig_image.size - max_pixels = int(max_megapixels * 1_048_576) - orig_pixels = width * height - patch_size = 32 # 保证分辨率为32的倍数 - # 动态调整分辨率 - if orig_pixels > max_pixels: - scale = (max_pixels / orig_pixels) ** 0.5 - new_width = max(1, int(width * scale)) - new_height = max(1, int(height * scale)) + + if use_input_mask: + # 使用输入的mask + if i < mask.shape[0]: # 确保有对应的mask + input_mask = mask[i:i+1] # 取对应的mask + else: + # 如果mask数量不足，使用最后一个mask + input_mask = mask[-1:] + + # 将mask转换为PIL图像并调整大小 + mask_pil = tensor2pil(input_mask) + if mask_pil.mode != 'L': + mask_pil = mask_pil.convert('L') + if mask_pil.size != orig_image.size: + mask_pil = mask_pil.resize(orig_image.size, Image.BILINEAR) + + # 翻转mask以确保正确的结果 + mask_pil = ImageOps.invert(mask_pil) + + # 直接使用输入mask + _mask = mask_pil + log(f"[BiRefNetUltra_RBS] 使用输入mask，跳过模型推理", message_type='info') else: - new_width, new_height = width, height - # 向下取整为patch_size的倍数 - new_width = (new_width // patch_size) * patch_size - new_height = (new_height // patch_size) * patch_size - new_width = max(patch_size, new_width) - new_height = max(patch_size, new_height) - inference_image_size = (new_width, new_height) - log(f"[BiRefNetUltraV3_RBS] 原始尺寸: {width}x{height}, 实际推理尺寸: {inference_image_size[0]}x{inference_image_size[1]}", message_type='info') - transform_image = transforms.Compose([ - transforms.Resize(inference_image_size), - transforms.ToTensor(), - transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) - ]) - inference_image = transform_image(orig_image).unsqueeze(0).to(device) - with torch.no_grad(): - preds = birefnet_model(inference_image)[-1].sigmoid().cpu() - pred = preds[0].squeeze() - pred_pil = transforms.ToPILImage()(pred) - # 先还原到inference_image_size，再还原到原图尺寸 - _mask = pred_pil.resize(inference_image_size, Image.BILINEAR) - _mask = _mask.resize(orig_image.size, Image.BILINEAR) - brightness_image = ImageEnhance.Brightness(_mask) - _mask = brightness_image.enhance(factor=1.08) - _mask = image2mask(_mask) + # 正常进行模型推理 + max_pixels = int(max_megapixels * 1_048_576) + orig_pixels = width * height + patch_size = 32 # 保证分辨率为32的倍数 + # 动态调整分辨率 + if orig_pixels > max_pixels: + scale = (max_pixels / orig_pixels) ** 0.5 + new_width = max(1, int(width * scale)) + new_height = max(1, int(height * scale)) + else: + new_width, new_height = width, height + # 向下取整为patch_size的倍数 + new_width = (new_width // patch_size) * patch_size + new_height = (new_height // patch_size) * patch_size + new_width = max(patch_size, new_width) + new_height = max(patch_size, new_height) + inference_image_size = (new_width, new_height) + log(f"[BiRefNetUltra_RBS] 原始尺寸: {width}x{height}, 实际推理尺寸: {inference_image_size[0]}x{inference_image_size[1]}", message_type='info') + transform_image = transforms.Compose([ + transforms.Resize(inference_image_size), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + inference_image = transform_image(orig_image).unsqueeze(0).to(device) + with torch.no_grad(): + preds = birefnet_model(inference_image)[-1].sigmoid().cpu() + pred = preds[0].squeeze() + pred_pil = transforms.ToPILImage()(pred) + # 先还原到inference_image_size，再还原到原图尺寸 + _mask = pred_pil.resize(inference_image_size, Image.BILINEAR) + _mask = _mask.resize(orig_image.size, Image.BILINEAR) + brightness_image = ImageEnhance.Brightness(_mask) + _mask = brightness_image.enhance(factor=1.08) + + # 统一处理mask格式 if not isinstance(_mask, Image.Image): _mask = tensor2pil(_mask) if _mask.mode != 'L': _mask = _mask.convert('L') if _mask.size != orig_image.size: _mask = _mask.resize(orig_image.size, Image.BILINEAR) + ret_image = RGB2RGBA(orig_image, _mask) ret_images.append(pil2tensor(ret_image)) ret_masks.append(image2mask(_mask)) comfy_pbar.update(1) tqdm_pbar.update(1) + log(f"{self.NODE_NAME} Processed {len(ret_masks)} image(s).", message_type='finish') return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),) @@ -328,13 +437,13 @@ class ProcessDetails_RBS: # 节点注册映射 NODE_CLASS_MAPPINGS = { "TransparentBackgroundUltra_RBS": TransparentBackgroundUltra_RBS, - "BiRefNetUltraV3_RBS": BiRefNetUltraV3_RBS, + "BiRefNetUltra_RBS": BiRefNetUltra_RBS, "MaskProcessDetails_RBS": ProcessDetails_RBS } NODE_DISPLAY_NAME_MAPPINGS = { "TransparentBackgroundUltra_RBS": "Transparent Background Ultra (RBS)", - "BiRefNetUltraV3_RBS": "BiRefNet Ultra V3 (RBS)", + "BiRefNetUltra_RBS": "BiRefNet Ultra (RBS)", "MaskProcessDetails_RBS": "Mask Process Details (RBS)" } diff --git a/pyproject.toml b/pyproject.toml index 864d6b1..0b68b73 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "removebackgroundsuite" description = "A matting toolkit based on ComfyUI, supporting multiple matting models and detail processing methods." -version = "1.1.2" +version = "1.2.0" license = {file = "LICENSE"} dependencies = ["torch", "numpy", "Pillow", "torchvision", "opencv-python", "scipy", "transformers", "tqdm", "pymatting", "guided-filter", "vitmatte", "timm"] diff --git a/rules.md b/rules.md index 697a884..0729e60 100644 --- a/rules.md +++ b/rules.md @@ -16,7 +16,7 @@ ComfyUI-RemoveBackgroundSuite 是一个基于 ComfyUI 的抠图工具包，支 2. **文件命名规范** - Python 文件使用小写字母和下划线：`model_utils.py` - - 类名使用 PascalCase：`BiRefNetUltraV3_RBS` + - 类名使用 PascalCase：`BiRefNetUltra_RBS` - 函数和变量使用 snake_case：`process_mask` 3. **注释规范**