125 lines
3.5 KiB
Python
125 lines
3.5 KiB
Python
import os
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from torchvision import transforms
|
|
from torch.autograd import Variable
|
|
|
|
from .M3Net import M3Net
|
|
from .data.custom_transforms import *
|
|
|
|
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)
|
|
|
|
|
|
class M3Net_ModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"m3net": (
|
|
[
|
|
"M3Net-R", "M3Net-S"
|
|
], {
|
|
"default": "M3Net-S"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("M3NET",)
|
|
RETURN_NAMES = ("m3net",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "m3net"
|
|
|
|
def load_model(self, m3net):
|
|
if m3net == "M3Net-S":
|
|
model = M3Net(embed_dim=512, dim=64, img_size=384, method="M3Net-S")
|
|
else:
|
|
model = M3Net(embed_dim=384, dim=64, img_size=384, method="M3Net-R")
|
|
model_path = os.path.join(current_directory, "weights/M3Net-S.pth")
|
|
model.load_state_dict(torch.load(model_path, map_location=device))
|
|
model.to(device)
|
|
model.eval()
|
|
return (model,)
|
|
|
|
|
|
class M3Net_Interface:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"m3net": ("M3NET",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK",)
|
|
RETURN_NAMES = ("image", "mask",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "m3net"
|
|
|
|
def process(self, m3net, image):
|
|
processed_images = []
|
|
processed_masks = []
|
|
|
|
# import pdb;pdb.set_trace()
|
|
for image in image:
|
|
image_pil = tensor2pil(image).convert("RGB")
|
|
comp = []
|
|
comp.append(static_resize(size=[384, 384]))
|
|
comp.append(tonumpy())
|
|
comp.append(normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]))
|
|
comp.append(totensor())
|
|
sample = {"image": image_pil, "shape": image_pil.size}
|
|
data_batch = transforms.Compose(comp)(sample)
|
|
|
|
images = data_batch['image']
|
|
image_w, image_h = data_batch['shape']
|
|
image_w, image_h = int(image_w), int(image_h)
|
|
images = Variable(images.cuda()[None, :, :, :])
|
|
|
|
outputs_saliency = m3net(images)
|
|
mask_1_1 = outputs_saliency[-1]
|
|
pred = torch.sigmoid(mask_1_1)
|
|
|
|
transform = transforms.Compose([
|
|
transforms.ToPILImage(),
|
|
transforms.Resize((image_h, image_w))
|
|
])
|
|
|
|
pred = pred.squeeze(0)
|
|
pred = transform(pred)
|
|
|
|
no_bg_image = Image.new("RGBA", pred.size, (0, 0, 0, 0))
|
|
no_bg_image.paste(image_pil, mask=pred)
|
|
|
|
new_im_tensor = pil2tensor(no_bg_image)
|
|
pil_im_tensor = pil2tensor(pred)
|
|
|
|
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
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"M3Net_ModelLoader": M3Net_ModelLoader,
|
|
"M3Net_Interface": M3Net_Interface
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"M3Net_ModelLoader": "M3Net Model Loader",
|
|
"M3Net_Interface": "M3Net Interface"
|
|
}
|