71 lines
2.4 KiB
Python
71 lines
2.4 KiB
Python
from transparent_background import Remover
|
|
import torchvision.transforms.v2 as T
|
|
from PIL import Image
|
|
import torch
|
|
import numpy as np
|
|
|
|
class TransparentBackgroundRembg:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", ),
|
|
"model_path": ("STRING", {"default": "/kaggle/working/latest.pth", "multiline": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
|
|
RETURN_NAMES = ("Foreground", "Mask", "Foreground_Transparent")
|
|
FUNCTION = "remove_background"
|
|
CATEGORY = "image"
|
|
|
|
def remove_background(self, image: torch.Tensor, model_path: str):
|
|
remover = Remover(mode='base-nightly', ckpt=model_path, device='cuda:0')
|
|
bgrgba = None
|
|
|
|
image = image.permute([0, 3, 1, 2])
|
|
output = []
|
|
for img in image:
|
|
img = T.ToPILImage()(img)
|
|
img = remover.process(img, type='rgba')
|
|
output.append(T.ToTensor()(img))
|
|
|
|
output = torch.stack(output, dim=0)
|
|
output = output.permute([0, 2, 3, 1])
|
|
# mask = output[:, :, :, 3] if output.shape[3] == 4 else torch.ones_like(output[:, :, :, 0])
|
|
|
|
# if output.shape[3] == 4:
|
|
# print("Has 4 channels")
|
|
# # Extract the RGB channels and the alpha channel separately
|
|
# output_rgb = output[:, :, :, :3]
|
|
# mask = output[:, :, :, 3]
|
|
# else:
|
|
# print("Not having 4 channels")
|
|
# print("Has: {}".format(output.shape[3]))
|
|
# # If there's no alpha channel, just use the output as-is
|
|
# output_rgb = output
|
|
# mask = torch.ones_like(output[:, :, :, 0]) # Create a mask of ones
|
|
# return(output_rgb, mask,)
|
|
|
|
|
|
if output.shape[3] == 4:
|
|
print("Has 4 channels")
|
|
alpha_channel = output[:, :, :, 3]
|
|
mask = alpha_channel
|
|
output_rgb = output[:, :, :, :3]
|
|
# output_rgb[~mask] = 0 # Set the background to black where alpha was 0
|
|
else:
|
|
print("Not having 4 channels")
|
|
print("Has: {}".format(output.shape[3]))
|
|
output_rgb = output
|
|
mask = torch.ones_like(output[:, :, :, 0]) # Create a mask of ones
|
|
|
|
return output_rgb, mask, output
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Transparentbackground RemBg": TransparentBackgroundRembg
|
|
}
|