CavinHuang-comfyui-nodes-docs/docs/BiRefNet_Remove_Background.md

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

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)