feat: 升级到v1.2.0 - 添加mask输入功能和智能遮罩检测
- 为BiRefNet Ultra和Transparent Background Ultra节点添加mask输入 - 新增'自动适用原图遮罩'参数，支持智能检测和使用输入的mask - 优化节点名称，移除BiRefNetUltraV3_RBS中的'V3'后缀 - 修复输入mask应用时的翻转问题，确保输出结果正确 - 改进批量处理逻辑，支持图像和mask的自动匹配 - 更新版本号到v1.2.0
This commit is contained in:
Cyber Dick Lang
@@ -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)"
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user