| import os |
| import io |
| import torch |
| import numpy as np |
| from PIL import Image |
| import cv2 |
| import pytorch_lightning as pl |
| from model import ISNetDIS, ISNetGTEncoder, U2NET, U2NET_full2, U2NET_lite2, MODNet |
|
|
| model = None |
|
|
| def get_mask(model, input_img): |
| h, w = input_img.shape[0], input_img.shape[1] |
| ph, pw = 0, 0 |
| tmpImg = np.zeros([h, w, 3], dtype=np.float16) |
| 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(): |
| 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)), (w, h))[:, :, np.newaxis] |
| return pred |
|
|
| def get_net(net_name): |
| if net_name == "isnet": |
| return ISNetDIS() |
| elif net_name == "isnet_is": |
| return ISNetDIS() |
| elif net_name == "isnet_gt": |
| return ISNetGTEncoder() |
| elif net_name == "u2net": |
| return U2NET_full2() |
| elif net_name == "u2netl": |
| return U2NET_lite2() |
| elif net_name == "modnet": |
| return MODNet() |
| raise NotImplemented |
|
|
| |
| class AnimeSegmentation(pl.LightningModule): |
| def __init__(self, net_name): |
| super().__init__() |
| assert net_name in ["isnet_is", "isnet", "isnet_gt", "u2net", "u2netl", "modnet"] |
| self.net = get_net(net_name) |
| if net_name == "isnet_is": |
| self.gt_encoder = get_net("isnet_gt") |
| for param in self.gt_encoder.parameters(): |
| param.requires_grad = False |
| else: |
| self.gt_encoder = None |
|
|
| @classmethod |
| def try_load(cls, net_name, ckpt_path, map_location=None): |
| state_dict = torch.load(ckpt_path, map_location=map_location) |
| if "epoch" in state_dict: |
| return cls.load_from_checkpoint(ckpt_path, net_name=net_name, map_location=map_location) |
| else: |
| model = cls(net_name) |
| if any([k.startswith("net.") for k, v in state_dict.items()]): |
| model.load_state_dict(state_dict) |
| else: |
| model.net.load_state_dict(state_dict) |
| return model |
|
|
| def forward(self, x): |
| if isinstance(self.net, ISNetDIS): |
| return self.net(x)[0][0].sigmoid() |
| if isinstance(self.net, ISNetGTEncoder): |
| return self.net(x)[0][0].sigmoid() |
| elif isinstance(self.net, U2NET): |
| return self.net(x)[0].sigmoid() |
| elif isinstance(self.net, MODNet): |
| return self.net(x, True)[2] |
| raise NotImplemented |
|
|
| def load_model(): |
| global model |
|
|
| if torch.cuda.is_available(): |
| device = 'cuda' |
| else: |
| device = 'cpu' |
| |
| model = AnimeSegmentation.try_load('isnet_is', 'anime-seg/isnetis.ckpt', device) |
| model.eval() |
| model.to(device) |
| |
| def animeseg(image): |
| global model |
|
|
| if not image: |
| return None |
|
|
| if not model: |
| model = load_model() |
|
|
| img = np.array(image, dtype=np.uint8) |
| mask = get_mask(model, img) |
| img = np.concatenate((mask * img + 1 - mask, mask * 255), axis=2).astype(np.uint8) |
| return img |
|
|
| def pil_to_webp(img): |
| buffer = io.BytesIO() |
| img.save(buffer, 'webp') |
|
|
| return buffer.getvalue() |
|
|
| def bin_to_base64(bin): |
| return base64.b64encode(bin).decode('ascii') |
|
|
