add GetMaskByBiRefNet and BlurFusionForegroundEstimation

2025-01-04 13:15:04 +08:00
parent 9b7262e588
commit 7570a99ff3
4 changed files with 137 additions and 52 deletions
@@ -164,7 +164,135 @@ class LoadRembgByBiRefNetModel:
        return [(biRefNet_model, version)]


-class RembgByBiRefNetAdvanced:
+class GetMaskByBiRefNet:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("BiRefNetMODEL",),
+                "images": ("IMAGE",),
+                "width": ("INT",
+                          {
+                              "default": 1024,
+                              "min": 0,
+                              "max": 16384,
+                              "tooltip": "The width of the preprocessed image, does not affect the final output image size"
+                          }),
+                "height": ("INT",
+                           {
+                               "default": 1024,
+                               "min": 0,
+                               "max": 16384,
+                               "tooltip": "The height of the preprocessed image, does not affect the final output image size"
+                           }),
+                "upscale_method": (["bislerp", "nearest-exact", "bilinear", "area", "bicubic"],
+                                   {
+                                       "default": "bilinear",
+                                       "tooltip": "Interpolation method for post-processing mask"
+                                   }),
+                "mask_threshold": ("FLOAT", {"default": 0.000, "min": 0.0, "max": 1.0, "step": 0.004, }),
+            }
+        }
+
+    RETURN_TYPES = ("MASK",)
+    RETURN_NAMES = ("mask",)
+    FUNCTION = "get_mask"
+    CATEGORY = "rembg/BiRefNet"
+
+    def get_mask(self, model, images, width=1024, height=1024, upscale_method='bilinear', mask_threshold=0.000):
+        model, version = model
+        model_device_type = next(model.parameters()).device.type
+        b, h, w, c = images.shape
+        image_bchw = images.permute(0, 3, 1, 2)
+
+        image_preproc = ImagePreprocessor(resolution=(height, width))
+        if VERSION[0] == version:
+            im_tensor = image_preproc.old_proc(image_bchw)
+        else:
+            im_tensor = image_preproc.proc(image_bchw)
+
+        _mask_bchw = []
+        for each_image in im_tensor:
+            with torch.no_grad():
+                each_mask = model(each_image.unsqueeze(0).to(model_device_type))[-1].sigmoid().cpu()
+            _mask_bchw.append(each_mask)
+            del each_mask
+
+        mask_bchw = torch.cat(_mask_bchw, dim=0)
+        del _mask_bchw
+        # 遮罩大小需还原为与原图一致
+        mask = comfy.utils.common_upscale(mask_bchw, w, h, upscale_method, "disabled")
+        # (b, 1, h, w)
+        if mask_threshold > 0:
+            mask = filter_mask(mask, threshold=mask_threshold)
+        # else:
+        #   似乎几乎无影响
+        #     mask = normalize_mask(mask)
+
+        return mask.squeeze(1),
+
+
+class BlurFusionForegroundEstimation:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "masks": ("MASK",),
+                "blur_size": ("INT", {"default": 91, "min": 1, "max": 255, "step": 2, }),
+                "blur_size_two": ("INT", {"default": 7, "min": 1, "max": 255, "step": 2, }),
+                "fill_color": ("BOOLEAN", {"default": False}),
+                "color": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFF, "step": 1, "display": "color"}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE", "MASK",)
+    RETURN_NAMES = ("image", "mask",)
+    FUNCTION = "get_foreground"
+    CATEGORY = "rembg/BiRefNet"
+    DESCRIPTION = "Approximate Fast Foreground Colour Estimation. https://github.com/Photoroom/fast-foreground-estimation"
+
+    def get_foreground(self, images, masks, blur_size=91, blur_size_two=7, fill_color=False, color=None):
+        b, h, w, c = images.shape
+        if b != masks.shape[0]:
+            raise ValueError("images and masks must have the same batch size")
+
+        image_bchw = images.permute(0, 3, 1, 2)
+
+        if masks.dim() == 3:
+            # (b, h, w) => (b, 1, h, w)
+            out_masks = masks.unsqueeze(1)
+
+        # (b, c, h, w)
+        _image_masked = refine_foreground(image_bchw, out_masks, r1=blur_size, r2=blur_size_two)
+        # (b, c, h, w) => (b, h, w, c)
+        _image_masked = _image_masked.permute(0, 2, 3, 1)
+        if fill_color and color is not None:
+            r = torch.full([b, h, w, 1], ((color >> 16) & 0xFF) / 0xFF)
+            g = torch.full([b, h, w, 1], ((color >> 8) & 0xFF) / 0xFF)
+            b = torch.full([b, h, w, 1], (color & 0xFF) / 0xFF)
+            # (b, h, w, 3)
+            background_color = torch.cat((r, g, b), dim=-1)
+            # (b, 1, h, w) => (b, h, w, 3)
+            apply_mask = out_masks.permute(0, 2, 3, 1).expand_as(_image_masked)
+            out_images = _image_masked * apply_mask + background_color * (1 - apply_mask)
+            # (b, h, w, 3)=>(b, h, w, 3)
+            del background_color, apply_mask
+            out_masks = out_masks.squeeze(1)
+        else:
+            # (b, 1, h, w) => (b, h, w)
+            out_masks = out_masks.squeeze(1)
+            # image的非mask对应部分设为透明 => (b, h, w, 4)
+            out_images = add_mask_as_alpha(_image_masked.cpu(), out_masks.cpu())
+
+        del _image_masked
+
+        return out_images, out_masks
+
+
+class RembgByBiRefNetAdvanced(GetMaskByBiRefNet, BlurFusionForegroundEstimation):

    @classmethod
    def INPUT_TYPES(cls):
@@ -205,57 +333,10 @@ class RembgByBiRefNetAdvanced:
    CATEGORY = "rembg/BiRefNet"

    def rem_bg(self, model, images, upscale_method='bilinear', width=1024, height=1024, blur_size=91, blur_size_two=7, fill_color=False, color=None, mask_threshold=0.000):
-        model, version = model
-        model_device_type = next(model.parameters()).device.type
-        b, h, w, c = images.shape
-        image_bchw = images.permute(0, 3, 1, 2)

-        image_preproc = ImagePreprocessor(resolution=(1024, 1024))
-        if VERSION[0] == version:
-            im_tensor = image_preproc.old_proc(image_bchw)
-        else:
-            im_tensor = image_preproc.proc(image_bchw)
+        masks = super().get_mask(model, images, width, height, upscale_method, mask_threshold)

-        _mask_bchw = []
-        for each_image in im_tensor:
-            with torch.no_grad():
-                each_mask = model(each_image.unsqueeze(0).to(model_device_type))[-1].sigmoid().cpu()
-            _mask_bchw.append(each_mask)
-            del each_mask
-
-        mask_bchw = torch.cat(_mask_bchw, dim=0)
-        del _mask_bchw
-        # 遮罩大小需还原为与原图一致
-        mask = comfy.utils.common_upscale(mask_bchw, w, h, 'bilinear', "disabled")
-        # (b, 1, h, w)
-        if mask_threshold > 0:
-            out_masks = filter_mask(mask, threshold=mask_threshold)
-        else:
-            out_masks = normalize_mask(mask)
-
-        # (b, c, h, w)
-        _image_masked = refine_foreground(image_bchw, out_masks, r1=blur_size, r2=blur_size_two)
-        # (b, c, h, w) => (b, h, w, c)
-        _image_masked = _image_masked.permute(0, 2, 3, 1)
-        if fill_color and color is not None:
-            r = torch.full([b, h, w, 1], ((color >> 16) & 0xFF) / 0xFF)
-            g = torch.full([b, h, w, 1], ((color >> 8) & 0xFF) / 0xFF)
-            b = torch.full([b, h, w, 1], (color & 0xFF) / 0xFF)
-            # (b, h, w, 3)
-            background_color = torch.cat((r, g, b), dim=-1)
-            # (b, 1, h, w) => (b, h, w, 3)
-            apply_mask = out_masks.permute(0, 2, 3, 1).expand_as(_image_masked)
-            out_images = _image_masked * apply_mask + background_color * (1 - apply_mask)
-            # (b, h, w, 3)=>(b, h, w, 3)
-            del background_color, apply_mask
-            out_masks = out_masks.squeeze(1)
-        else:
-            # (b, 1, h, w) => (b, h, w)
-            out_masks = out_masks.squeeze(1)
-            # image的非mask对应部分设为透明 => (b, h, w, 4)
-            out_images = add_mask_as_alpha(_image_masked.cpu(), out_masks.cpu())
-
-        del _image_masked
+        out_images, out_masks = super().get_foreground(images, masks=masks[0], blur_size=blur_size, blur_size_two=blur_size_two, fill_color=fill_color, color=color)

        return out_images, out_masks

@@ -285,6 +366,8 @@ NODE_CLASS_MAPPINGS = {
    "LoadRembgByBiRefNetModel": LoadRembgByBiRefNetModel,
    "RembgByBiRefNet": RembgByBiRefNet,
    "RembgByBiRefNetAdvanced": RembgByBiRefNetAdvanced,
+    "GetMaskByBiRefNet": GetMaskByBiRefNet,
+    "BlurFusionForegroundEstimation": BlurFusionForegroundEstimation,
 }

 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -292,4 +375,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
    "LoadRembgByBiRefNetModel": "LoadRembgByBiRefNetModel",
    "RembgByBiRefNet": "RembgByBiRefNet",
    "RembgByBiRefNetAdvanced": "RembgByBiRefNetAdvanced",
+    "GetMaskByBiRefNet": "GetMaskByBiRefNet",
+    "BlurFusionForegroundEstimation": "BlurFusionForegroundEstimation",
 }
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui_birefnet_ll"
-description = "Sync with version of BiRefNet. NODES:AutoDownloadBiRefNetModel, LoadRembgByBiRefNetModel, RembgByBiRefNet."
-version = "1.0.7"
+description = "Sync with version of BiRefNet. NODES:AutoDownloadBiRefNetModel, LoadRembgByBiRefNetModel, RembgByBiRefNet, RembgByBiRefNetAdvanced, GetMaskByBiRefNet, BlurFusionForegroundEstimation."
+version = "1.0.8"
 license = {file = "LICENSE"}
 dependencies = ["numpy", "opencv-python", "timm"]