1.9 KiB
1.9 KiB
tags
| tags | ||
|---|---|---|
|
🔧 Image Remove Background
Documentation
- Class name:
ImageRemoveBackground+ - Category:
essentials - Output node:
False
This node is designed to remove the background from images, utilizing the rembg library to process each image in a batch and generate a corresponding mask for each image.
Input types
Required
rembg_session- The session object for the rembg library, required for processing the images to remove their backgrounds.
- Comfy dtype:
REMBG_SESSION - Python dtype:
object
image- The input image or batch of images to have their backgrounds removed.
- Comfy dtype:
IMAGE - Python dtype:
torch.tensor
Output types
image- Comfy dtype:
IMAGE - The batch of images with their backgrounds removed.
- Python dtype:
torch.tensor
- Comfy dtype:
mask- Comfy dtype:
MASK - A mask or batch of masks corresponding to the removed backgrounds of the input images.
- Python dtype:
torch.tensor
- Comfy dtype:
Usage tips
- Infra type:
GPU - Common nodes: unknown
Source code
class ImageRemoveBackground:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"rembg_session": ("REMBG_SESSION",),
"image": ("IMAGE",),
},
}
RETURN_TYPES = ("IMAGE", "MASK",)
FUNCTION = "execute"
CATEGORY = "essentials"
def execute(self, rembg_session, image):
from rembg import remove as rembg
image = p(image)
output = []
for img in image:
img = T.ToPILImage()(img)
img = rembg(img, session=rembg_session)
output.append(T.ToTensor()(img))
output = torch.stack(output, dim=0)
output = pb(output)
mask = output[:, :, :, 3] if output.shape[3] == 4 else torch.ones_like(output[:, :, :, 0])
return(output[:, :, :, :3], mask,)