From 0ab09df6c454e324c676ec6cc34041e705f67403 Mon Sep 17 00:00:00 2001
From: yolain <yolain@163.com>
Date: Thu, 19 Dec 2024 22:05:20 +0800
Subject: [PATCH] Add `add_background` of widget on `easy imageRembg`

---
 py/image.py      | 16 +++++++++++++---
 py/libs/image.py |  5 +++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/py/image.py b/py/image.py
index fa64ac7..d85abf6 100644
--- a/py/image.py
+++ b/py/image.py
@@ -1,5 +1,6 @@
 import os
 import json
+import copy
 import hashlib
 import folder_paths
 import torch
@@ -17,7 +18,7 @@ from torchvision.transforms.functional import to_pil_image
 from .libs.log import log_node_info
 from .libs.utils import AlwaysEqualProxy, ByPassTypeTuple
 from .libs.cache import cache, update_cache, remove_cache
-from .libs.image import pil2tensor, tensor2pil, ResizeMode, get_new_bounds, RGB2RGBA, image2mask
+from .libs.image import pil2tensor, tensor2pil, ResizeMode, get_new_bounds, RGB2RGBA, image2mask, empty_image
 from .libs.colorfix import adain_color_fix, wavelet_color_fix
 from .libs.chooser import ChooserMessage, ChooserCancelled
 from .config import REMBG_DIR, REMBG_MODELS, HUMANPARSING_MODELS, MEDIAPIPE_MODELS, MEDIAPIPE_DIR
@@ -801,10 +802,10 @@ class imageRemBg:
         "rem_mode": (("RMBG-2.0", "RMBG-1.4","Inspyrenet"), {"default": "RMBG-1.4"}),
         "image_output": (["Hide", "Preview", "Save", "Hide/Save"], {"default": "Preview"}),
         "save_prefix": ("STRING", {"default": "ComfyUI"}),
-
       },
       "optional":{
         "torchscript_jit": ("BOOLEAN", {"default": False}),
+        "add_background": (["none", "white", "black"], {"default": "none"})
       },
       "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
     }
@@ -816,7 +817,8 @@ class imageRemBg:
 
   CATEGORY = "EasyUse/Image"
 
-  def remove(self, rem_mode, images, image_output, save_prefix, torchscript_jit=False, prompt=None, extra_pnginfo=None):
+
+  def remove(self, rem_mode, images, image_output, save_prefix, torchscript_jit=False, add_background='none',prompt=None, extra_pnginfo=None):
     new_images = list()
     masks = list()
     if rem_mode == "RMBG-2.0":
@@ -911,6 +913,14 @@ class imageRemBg:
       new_images = torch.cat(new_images, dim=0)
       masks = torch.cat(masks, dim=0)
 
+    if add_background != 'none':
+
+      _layer = tensor2pil(new_images)
+      _canvas = Image.new('RGBA', _layer.size, (255,255,255,255) if add_background == 'white' else (0, 0, 0, 255))
+      _canvas.paste(_layer, mask=_layer)
+      new_images = pil2tensor(_canvas)
+
+
     results = easySave(new_images, save_prefix, image_output, prompt, extra_pnginfo)
 
     if image_output in ("Hide", "Hide/Save"):
diff --git a/py/libs/image.py b/py/libs/image.py
index c4ec189..5c3efe9 100644
--- a/py/libs/image.py
+++ b/py/libs/image.py
@@ -105,6 +105,11 @@ class blendImage:
     return blended_image
 
 
+def empty_image(width, height, batch_size=1, color=0):
+    r = torch.full([batch_size, height, width, 1], ((color >> 16) & 0xFF) / 0xFF)
+    g = torch.full([batch_size, height, width, 1], ((color >> 8) & 0xFF) / 0xFF)
+    b = torch.full([batch_size, height, width, 1], ((color) & 0xFF) / 0xFF)
+    return torch.cat((r, g, b), dim=-1)
 
 
 class ResizeMode(Enum):