BRIA_RMBG.py

import torch, os
import torch.nn.functional as F
from PIL import Image
from .briarmbg import BriaRMBG
from torchvision.transforms.functional import normalize
import numpy as np

current_directory = os.path.dirname(os.path.abspath(__file__))
device = "cuda" if torch.cuda.is_available() else "cpu"

def tensor2pil(image):
    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))

def pil2tensor(image):
    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)

def resize_image(image):
    image = image.convert('RGB')
    model_input_size = (1024, 1024)
    image = image.resize(model_input_size, Image.BILINEAR)
    return image


class BRIA_RMBG_ModelLoader:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
            }
        }

    RETURN_TYPES = ("RMBGMODEL",)
    RETURN_NAMES = ("rmbgmodel",)
    FUNCTION = "load_model"
    CATEGORY = "BRIA RMBG"
  
    def load_model(self):
        net = BriaRMBG()
        model_path = os.path.join(current_directory, "RMBG-1.4/model.pth")
        net.load_state_dict(torch.load(model_path, map_location=device))
        net.to(device)
        net.eval() 
        return [net]


class BRIA_RMBG:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "rmbgmodel": ("RMBGMODEL",),
                "image": ("IMAGE",),
            }
        }

    RETURN_TYPES = ("IMAGE", "MASK", )
    RETURN_NAMES = ("image", "mask", )
    FUNCTION = "remove_background"
    CATEGORY = "AI API"
  
    def remove_background(self, rmbgmodel, image):
        processed_images = []
        processed_masks = []

        for image in image:
            orig_image = tensor2pil(image)
            w,h = orig_image.size
            image = resize_image(orig_image)
            im_np = np.array(image)
            im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1)
            im_tensor = torch.unsqueeze(im_tensor,0)
            im_tensor = torch.divide(im_tensor,255.0)
            im_tensor = normalize(im_tensor,[0.5,0.5,0.5],[1.0,1.0,1.0])
            if torch.cuda.is_available():
                im_tensor=im_tensor.cuda()

            result=rmbgmodel(im_tensor)
            result = torch.squeeze(F.interpolate(result[0][0], size=(h,w), mode='bilinear') ,0)
            ma = torch.max(result)
            mi = torch.min(result)
            result = (result-mi)/(ma-mi)    
            im_array = (result*255).cpu().data.numpy().astype(np.uint8)
            pil_im = Image.fromarray(np.squeeze(im_array))
            new_im = Image.new("RGBA", pil_im.size, (0,0,0,0))
            new_im.paste(orig_image, mask=pil_im)

            new_im_tensor = pil2tensor(new_im)  
            pil_im_tensor = pil2tensor(pil_im)  

            processed_images.append(new_im_tensor)
            processed_masks.append(pil_im_tensor)

        new_ims = torch.cat(processed_images, dim=0)
        new_masks = torch.cat(processed_masks, dim=0)

        return new_ims, new_masks