Merge pull request #519 from DenisKochetov/partcrafter-rembg

add white background

Author: MrForExample

Gen_3D_Modules/Hunyuan3D_2_1/hy3dshape/hy3dshape/rembg.py

@@ -15,13 +15,14 @@
 from PIL import Image
 from rembg import remove, new_session
 import numpy as np
+from typing import Optional, Tuple
 
 
 class BackgroundRemover():
     def __init__(self, model_name='u2net'):
         self.session = new_session(model_name)
 
-    def __call__(self, image: Image.Image):
+    def __call__(self, image: Image.Image, background_color: Optional[Tuple[int, int, int]] = None):
         if image.mode in ('RGBA', 'LA') or image.mode != 'RGB':
             background = Image.new('RGB', image.size, (255, 255, 255))
             if image.mode in ('RGBA', 'LA'):
@@ -34,5 +35,11 @@ def __call__(self, image: Image.Image):
 
         if output.mode != 'RGBA':
             output = output.convert('RGBA')
-            
-        return output
+
+        if background_color is not None:
+            # Add background color to the output image
+            filled = Image.new("RGB", output.size, background_color)
+            filled.paste(output, mask=output.split()[-1])
+            return filled
+
+        return output

nodes.py

@@ -5889,7 +5889,7 @@ def generate(self, partcrafter_pipe, image, num_parts, seed, num_tokens, num_inf
             rmbg_worker = BackgroundRemover_2_1()
             if pil_image.mode == "RGBA":
                 pil_image = pil_image.convert('RGB')
-            pil_image = rmbg_worker(pil_image)
+            pil_image = rmbg_worker(pil_image, background_color=(255, 255, 255))
             del rmbg_worker
 
         # Set sampling version