118 lines
4.1 KiB
Python
118 lines
4.1 KiB
Python
import copy
|
|
import sys
|
|
import os
|
|
import torch
|
|
import cv2
|
|
import numpy as np
|
|
from torch.cuda import amp
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "../comfy_extras/anime_segmentation")) # noqa
|
|
from comfy_extras.anime_segmentation.train import AnimeSegmentation
|
|
|
|
|
|
class Segment:
|
|
net = ["isnet_is", "isnet", "u2net", "u2netl", "modnet"]
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"images": ("IMAGE", ),
|
|
"net": (s.net, {"default": "isnet_is"}),
|
|
"ckpt": ("STRING", {"default": "comfy_extras/anime_segmentation/saved_models/isnetis.ckpt"}),
|
|
"device": ("STRING", {"default": "cuda:0"}),
|
|
"precision": (["fp16", "fp32"], {"default": "fp16"}),
|
|
"image_size": ("INT", {"default": 1024}),
|
|
"remove": (["character", "background"], {"default": "background"}),
|
|
"alpha": (["yes", "no"], {"default": "yes"}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
FUNCTION = "segment"
|
|
|
|
CATEGORY = "image"
|
|
|
|
def get_mask(self, model, input_img, use_amp=True, s=640):
|
|
h0, w0 = h, w = input_img.shape[0], input_img.shape[1]
|
|
if h > w:
|
|
h, w = s, int(s * w / h)
|
|
else:
|
|
h, w = int(s * h / w), s
|
|
ph, pw = s - h, s - w
|
|
tmpImg = np.zeros([s, s, 3], dtype=np.float32)
|
|
tmpImg[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w] = cv2.resize(input_img, (w, h)) / 255
|
|
tmpImg = tmpImg.transpose((2, 0, 1))
|
|
tmpImg = torch.from_numpy(tmpImg).unsqueeze(0).type(torch.FloatTensor).to(model.device)
|
|
with torch.no_grad():
|
|
if use_amp:
|
|
with amp.autocast():
|
|
pred = model(tmpImg)
|
|
pred = pred.to(dtype=torch.float32)
|
|
else:
|
|
pred = model(tmpImg)
|
|
pred = pred[0, :, ph // 2:ph // 2 + h, pw // 2:pw // 2 + w]
|
|
pred = cv2.resize(pred.cpu().numpy().transpose((1, 2, 0)), (w0, h0))[:, :, np.newaxis]
|
|
return pred
|
|
|
|
def segment(self, images, net, ckpt, device, precision, image_size, remove, alpha):
|
|
ckpt = os.path.realpath(ckpt)
|
|
|
|
device = torch.device(device)
|
|
model = AnimeSegmentation.try_load(net, ckpt, device)
|
|
model.eval()
|
|
model.to(device)
|
|
|
|
outputs = []
|
|
first_mask = None
|
|
for image in images:
|
|
i = 255. * image.cpu().numpy()
|
|
img = cv2.cvtColor(np.clip(i, 0, 255).astype(np.uint8), cv2.COLOR_BGR2RGB)
|
|
mask = self.get_mask(model, img, use_amp=precision == "fp16", s=image_size)
|
|
|
|
img = np.concatenate(
|
|
(mask * img + 1 - mask, mask * 255), axis=2).astype(np.uint8)
|
|
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA)
|
|
|
|
if remove == "background":
|
|
# Store mask if we are first
|
|
if first_mask is None:
|
|
mask = img[:, :, 3].astype(np.float32) / 255.0
|
|
mask = torch.from_numpy(mask)
|
|
first_mask = 1. - mask
|
|
else:
|
|
# Copy the source image
|
|
i = i.copy()
|
|
|
|
# Add alpha channel if it isnt present
|
|
if i.shape[2] == 3:
|
|
i = np.concatenate([i, np.ones((i.shape[0], i.shape[1], 1), dtype=i.dtype) * 255], axis=2)
|
|
|
|
# Remove non transparent pixels in segmented image from source
|
|
mask = img[:, :, 3] > 0
|
|
i[mask] = [0, 0, 0, 0]
|
|
img = i
|
|
|
|
# Store mask if we are first
|
|
if first_mask is None:
|
|
mask = img[:, :, 3].astype(np.float32, copy=True) / 255.0
|
|
mask = torch.from_numpy(mask)
|
|
first_mask = 1. - mask
|
|
|
|
if alpha == "no":
|
|
# Remove alpha channel
|
|
img = img[:, :, :3]
|
|
|
|
img = img.astype(np.float32) / 255.0
|
|
img = torch.from_numpy(img)[None,]
|
|
outputs.append(img)
|
|
|
|
if len(outputs) == 1:
|
|
outputs = outputs[0]
|
|
|
|
return (outputs, first_mask,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AnimeSegmentation": Segment,
|
|
}
|