import os import sys import tqdm import gdown import torch import warnings import numpy as np import torch.nn.functional as F import torchvision.transforms as transforms from PIL import Image from packaging import version filepath = os.path.abspath(__file__) repopath = os.path.split(filepath)[0] sys.path.append(repopath) from InsPy.InSPyReNet import InSPyReNet_SwinB from InsPy.utils import * warnings.filterwarnings("ignore") CONFIG = { 'base': {'url': "https://ae-rt.oss-cn-beijing.aliyuncs.com/matting_module/inspy/ckpt_base.pth", 'md5': "d692e3dd5fa1b9658949d452bebf1cda", 'base_size': [1024, 1024], 'threshold': None, 'ckpt_name': "ckpt_base.pth", 'resize': dynamic_resize(L=1280)}, 'fast': {'url': "https://ae-rt.oss-cn-beijing.aliyuncs.com/matting_module/inspy/ckpt_fast.pth", 'md5': "9efdbfbcc49b79ef0f7891c83d2fd52f", 'base_size': [384, 384], 'threshold': 512, 'ckpt_name': "ckpt_fast.pth", 'resize': static_resize(size=[384, 384])} } class Remover: def __init__(self, fast=False, jit=False, device=None, ckpt=None): """ fast (default False): resize input into small size for fast computation jit (default False): use TorchScript for fast computation device (default cuda:0 if available): specifying device for computation ckpt (default download from server): specifying model checkpoint """ key = "fast" if fast else "base" self.meta = CONFIG[key] if device is not None: self.device=device else: self.device = "cpu" if torch.cuda.is_available(): self.device = "cuda:0" elif version.parse(torch.__version__) >= version.parse("1.13") and torch.backends.mps.is_available(): self.device = "mps:0" download = False if ckpt is None: ckpt_dir = os.path.expanduser(os.path.join('~', '.transparent-background')) if os.path.isdir(ckpt_dir) is False: os.makedirs(ckpt_dir, exist_ok=True) ckpt_name = self.meta['ckpt_name'] if not os.path.isfile(os.path.join(ckpt_dir, ckpt_name)): download = True elif self.meta['md5'] != hashlib.md5(open(os.path.join(ckpt_dir, ckpt_name), 'rb').read()).hexdigest(): download = True if download: gdown.download(self.meta['url'], os.path.join(ckpt_dir, ckpt_name)) else: ckpt_dir, ckpt_name = os.path.split(os.path.abspath(ckpt)) print(ckpt_dir, ckpt_name) self.model = InSPyReNet_SwinB(depth=64, pretrained=False, **self.meta) self.model.eval() self.model.load_state_dict(torch.load(os.path.join(ckpt_dir, ckpt_name), map_location='cpu'), strict=True) self.model = self.model.to(self.device) if jit: ckpt_name = self.meta['ckpt_name'].replace('.pth', '_{}.pt'.format(self.device)) try: traced_model = torch.jit.load(os.path.join(ckpt_dir, ckpt_name), map_location=self.device) del self.model self.model = traced_model except: traced_model = torch.jit.trace(self.model, torch.rand(1, 3, *self.meta['base_size']).to(self.device), strict=True) del self.model self.model = traced_model torch.jit.save(self.model, os.path.join(ckpt_dir, ckpt_name)) self.transform = transforms.Compose([static_resize(self.meta['base_size']) if jit else self.meta['resize'], tonumpy(), normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), totensor()]) self.background = None desc = 'Mode={}, Device={}, Torchscript={}'.format(key, self.device, 'enabled' if jit else 'disabled') # print('=' * (len(desc) + 2) + '\n', desc, '\n' + '=' * (len(desc) + 2)) print('Settings -> {}'.format(desc)) def process(self, img, type='rgba'): shape = img.size[::-1] x = self.transform(img) x = x.unsqueeze(0) x = x.to(self.device) with torch.no_grad(): pred = self.model(x) pred = F.interpolate(pred, shape, mode='bilinear', align_corners=True) pred = pred.data.cpu() pred = pred.numpy().squeeze() img = np.array(img) if type.startswith('['): type = [int(i) for i in type[1:-1].split(',')] if type == 'map': img = (np.stack([pred] * 3, axis=-1) * 255).astype(np.uint8) elif type == 'rgba': r, g, b = cv2.split(img) pred = (pred * 255).astype(np.uint8) img = cv2.merge([r, g, b, pred]) elif type == 'green': bg = np.stack([np.ones_like(pred)] * 3, axis=-1) * [120, 255, 155] img = img * pred[..., np.newaxis] + bg * (1 - pred[..., np.newaxis]) elif type == "white": bg = np.stack([np.ones_like(pred)] * 3, axis=-1) * [255, 255, 255] img = img * pred[..., np.newaxis] + bg * (1 - pred[..., np.newaxis]) elif len(type) == 3: bg = np.stack([np.ones_like(pred)] * 3, axis=-1) * type img = img * pred[..., np.newaxis] + bg * (1 - pred[..., np.newaxis]) elif type == 'blur': img = img * pred[..., np.newaxis] + cv2.GaussianBlur(img, (0, 0), 15) * (1 - pred[..., np.newaxis]) elif type == 'overlay': bg = (np.stack([np.ones_like(pred)] * 3, axis=-1) * [120, 255, 155] + img) // 2 img = bg * pred[..., np.newaxis] + img * (1 - pred[..., np.newaxis]) border = cv2.Canny(((pred > .5) * 255).astype(np.uint8), 50, 100) img[border != 0] = [120, 255, 155] elif type.lower().endswith(('.jpg', '.jpeg', '.png')): if self.background is None: self.background = cv2.cvtColor(cv2.imread(type), cv2.COLOR_BGR2RGB) self.background = cv2.resize(self.background, img.shape[:2][::-1]) img = img * pred[..., np.newaxis] + self.background * (1 - pred[..., np.newaxis]) if torch.cuda.is_available(): torch.cuda.empty_cache() return img.astype(np.uint8)