112 lines
4.9 KiB
Markdown
112 lines
4.9 KiB
Markdown
# Documentation
|
|
- Class name: BiRefNet_Remove_Background
|
|
- Category: BiRefNet🌟
|
|
- Output node: False
|
|
- Repo Ref: https://github.com/moon7star9/ComfyUI_BiRefNet_Universal
|
|
|
|
BiRefNet_Remove_Background节点是一个用于图像背景移除的处理节点。它接收已加载的BiRefNet模型和输入图像，通过深度学习方法准确分割前景和背景，并可以选择性地对前景进行优化处理。节点支持多种背景处理选项，包括透明背景和各种纯色背景，同时提供前景优化功能以提高边缘质量。
|
|
|
|
# Input types
|
|
## Required
|
|
- model
|
|
- 'model'参数接收从BiRefNet_Loader节点输出的模型实例，包含了预训练模型及其配置信息。
|
|
- Comfy dtype: BIREFNET_MODEL
|
|
- Python dtype: dict
|
|
- image
|
|
- 'image'参数是需要处理的输入图像。节点会自动调整图像大小以匹配模型的最佳工作分辨率。
|
|
- Comfy dtype: IMAGE
|
|
- Python dtype: torch.Tensor
|
|
- background_color
|
|
- 'background_color'参数允许用户选择输出图像的背景类型，可以是透明背景或指定的纯色背景（白色、黑色、绿色、蓝色、红色）。
|
|
- Comfy dtype: COMBO
|
|
- Python dtype: str
|
|
- use_refine
|
|
- 'use_refine'参数控制是否使用前景优化功能。启用此选项可以改善边缘质量，但可能会略微增加处理时间。
|
|
- Comfy dtype: BOOLEAN
|
|
- Python dtype: bool
|
|
|
|
# Output types
|
|
- IMAGE
|
|
- 输出处理后的图像，根据选择的背景类型可能是RGBA（透明背景）或RGB（纯色背景）格式。
|
|
- Comfy dtype: IMAGE
|
|
- Python dtype: torch.Tensor
|
|
- MASK
|
|
- 输出分割掩码，表示前景区域的二值图像。这个掩码可以用于其他节点的进一步处理。
|
|
- Comfy dtype: MASK
|
|
- Python dtype: torch.Tensor
|
|
|
|
# Usage tips
|
|
- Infra type: GPU/CPU
|
|
- 推荐在GPU上运行以获得更好的性能
|
|
- 使用透明背景选项时，输出将是RGBA格式
|
|
- 启用前景优化可以改善边缘质量，特别是对于复杂的边缘细节
|
|
- 输出的掩码可以用于其他图像处理节点的遮罩输入
|
|
|
|
# Source code
|
|
```python
|
|
class BiRefNet_Remove_Background:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("BIREFNET_MODEL",),
|
|
"image": ("IMAGE",),
|
|
"background_color": (["transparency"] + ["white", "black", "green", "blue", "red"], {"default": "transparency"}),
|
|
"use_refine": ("BOOLEAN", {"default": True})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
RETURN_NAMES = ("image", "mask")
|
|
FUNCTION = "inference"
|
|
CATEGORY = "BiRefNet🌟"
|
|
|
|
def inference(self, image, model, background_color, use_refine):
|
|
model_data = model
|
|
model = model_data["model"]
|
|
device = model_data["device"]
|
|
use_half_precision = model_data["half_precision"]
|
|
resolution = model_data["resolution"] # 直接使用模型的最佳分辨率
|
|
|
|
preprocessor = ImagePreprocessor(resolution)
|
|
processed_images = []
|
|
processed_masks = []
|
|
|
|
for img in image:
|
|
# 转换为PIL图像
|
|
orig_image = Image.fromarray(np.clip(255. * img.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
w, h = orig_image.size
|
|
|
|
# 预处理
|
|
image_tensor = preprocessor(orig_image.convert('RGB')).unsqueeze(0)
|
|
if use_half_precision:
|
|
image_tensor = image_tensor.half()
|
|
image_tensor = image_tensor.to(device)
|
|
|
|
# 推理
|
|
with torch.no_grad():
|
|
preds = model(image_tensor)[-1].sigmoid().cpu()
|
|
|
|
pred = preds[0].squeeze()
|
|
pred_pil = transforms.ToPILImage()(pred)
|
|
mask = pred_pil.resize((w, h))
|
|
|
|
|
|
# 选择应用前景优化
|
|
if use_refine:
|
|
refined_image = refine_foreground(orig_image, pred_pil, r=90) # 使用固定的r值即可（实验验证过调整r值对结果影响很小）
|
|
|
|
# 设置背景和前景
|
|
|
|
if background_color == "transparency":
|
|
result_image = Image.new("RGBA", (w, h), (0, 0, 0, 0))
|
|
result_image.paste(refined_image if use_refine else orig_image, mask=mask)
|
|
else:
|
|
result_image = Image.new("RGB", (w, h), background_color)
|
|
result_image.paste(refined_image if use_refine else orig_image, mask=mask)
|
|
|
|
# 转换回tensor
|
|
processed_images.append(torch.from_numpy(np.array(result_image).astype(np.float32) / 255.0).unsqueeze(0))
|
|
processed_masks.append(torch.from_numpy(np.array(mask).astype(np.float32) / 255.0).unsqueeze(0))
|
|
|
|
return torch.cat(processed_images, dim=0), torch.cat(processed_masks, dim=0) |