Add refine_foreground for ben2

locales/zh/nodeDefs.json

@@ -5532,6 +5532,12 @@
       },
       "image_output": {
         "name": "图像输出"
+      },
+      "add_background": {
+        "name": "添加背景"
+      },
+      "refine_foreground": {
+        "name": "优化前景"
       }
     },
     "outputs": {

py/nodes/image.py

@@ -828,7 +828,8 @@ class imageRemBg:
       },
       "optional":{
         "torchscript_jit": ("BOOLEAN", {"default": False}),
-        "add_background": (["none", "white", "black"], {"default": "none"})
+        "add_background": (["none", "white", "black"], {"default": "none"}),
+        "refine_foreground": ("BOOLEAN", {"default": False}),
       },
       "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
     }
@@ -841,7 +842,7 @@ class imageRemBg:
   CATEGORY = "EasyUse/Image"
 
 
-  def remove(self, rem_mode, images, image_output, save_prefix, torchscript_jit=False, add_background='none',prompt=None, extra_pnginfo=None):
+  def remove(self, rem_mode, images, image_output, save_prefix, torchscript_jit=False, add_background='none', refine_foreground=False, prompt=None, extra_pnginfo=None):
     new_images = list()
     masks = list()
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -941,7 +942,7 @@ class imageRemBg:
         if input_image.mode != 'RGBA':
           input_image = input_image.convert("RGBA")
 
-        mask, new_im = model.inference(input_image)
+        mask, new_im = model.inference(input_image, refine_foreground)
 
         new_im_tensor = pil2tensor(new_im)
         mask_tensor = pil2tensor(mask)