From ca83e199864befb7398e51cad3a1ff3cfde38d70 Mon Sep 17 00:00:00 2001 From: DenisKochetov Date: Thu, 7 Aug 2025 04:09:49 +0000 Subject: [PATCH] add white background --- .../Hunyuan3D_2_1/hy3dshape/hy3dshape/rembg.py | 13 ++++++++++--- nodes.py | 2 +- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/Gen_3D_Modules/Hunyuan3D_2_1/hy3dshape/hy3dshape/rembg.py b/Gen_3D_Modules/Hunyuan3D_2_1/hy3dshape/hy3dshape/rembg.py index 15aff70..408edb4 100644 --- a/Gen_3D_Modules/Hunyuan3D_2_1/hy3dshape/hy3dshape/rembg.py +++ b/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 @@ class BackgroundRemover(): 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 \ No newline at end of file diff --git a/nodes.py b/nodes.py index c5eddc9..e4b01fc 100644 --- a/nodes.py +++ b/nodes.py @@ -5889,7 +5889,7 @@ class PartCrafter_Generate: 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