addinng file

transparent_background/InSPyReNet.py

@@ -0,0 +1,152 @@
+import os
+import sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+filepath = os.path.abspath(__file__)
+repopath = os.path.split(filepath)[0]
+sys.path.append(repopath)
+
+from transparent_background.modules.layers import *
+from transparent_background.modules.context_module import *
+from transparent_background.modules.attention_module import *
+from transparent_background.modules.decoder_module import *
+
+from transparent_background.backbones.SwinTransformer import SwinB
+
+class InSPyReNet(nn.Module):
+    def __init__(self, backbone, in_channels, depth=64, base_size=[384, 384], threshold=512, **kwargs):
+        super(InSPyReNet, self).__init__()
+        self.backbone = backbone
+        self.in_channels = in_channels
+        self.depth = depth
+        self.base_size = base_size
+        self.threshold = threshold
+        
+        self.context1 = PAA_e(self.in_channels[0], self.depth, base_size=self.base_size, stage=0)
+        self.context2 = PAA_e(self.in_channels[1], self.depth, base_size=self.base_size, stage=1)
+        self.context3 = PAA_e(self.in_channels[2], self.depth, base_size=self.base_size, stage=2)
+        self.context4 = PAA_e(self.in_channels[3], self.depth, base_size=self.base_size, stage=3)
+        self.context5 = PAA_e(self.in_channels[4], self.depth, base_size=self.base_size, stage=4)
+
+        self.decoder = PAA_d(self.depth * 3, depth=self.depth, base_size=base_size, stage=2)
+
+        self.attention0 = SICA(self.depth    , depth=self.depth, base_size=self.base_size, stage=0, lmap_in=True)
+        self.attention1 = SICA(self.depth * 2, depth=self.depth, base_size=self.base_size, stage=1, lmap_in=True)
+        self.attention2 = SICA(self.depth * 2, depth=self.depth, base_size=self.base_size, stage=2              )
+
+        self.ret = lambda x, target: F.interpolate(x, size=target.shape[-2:], mode='bilinear', align_corners=False)
+        self.res = lambda x, size: F.interpolate(x, size=size, mode='bilinear', align_corners=False)
+        self.des = lambda x, size: F.interpolate(x, size=size, mode='nearest')
+        
+        self.image_pyramid = ImagePyramid(7, 1)
+        
+        self.transition0 = Transition(17)
+        self.transition1 = Transition(9)
+        self.transition2 = Transition(5)
+        
+        self.forward = self.forward_inference
+        
+    def to(self, device):
+        self.image_pyramid.to(device)
+        self.transition0.to(device)
+        self.transition1.to(device)
+        self.transition2.to(device)
+        super(InSPyReNet, self).to(device)
+        return self
+    
+    def cuda(self, idx=None):
+        if idx is None:
+            idx = torch.cuda.current_device()
+            
+        self.to(device="cuda:{}".format(idx))
+        return self
+
+    def eval(self):
+        super(InSPyReNet, self).train(False)
+        self.forward = self.forward_inference
+        return self
+    
+    def forward_inspyre(self, x):
+        B, _, H, W = x.shape
+    
+        x1, x2, x3, x4, x5 = self.backbone(x)
+        
+        x1 = self.context1(x1) #4
+        x2 = self.context2(x2) #4
+        x3 = self.context3(x3) #8
+        x4 = self.context4(x4) #16
+        x5 = self.context5(x5) #32
+
+        f3, d3 = self.decoder([x3, x4, x5]) #16
+
+        f3 = self.res(f3, (H // 4,  W // 4 ))
+        f2, p2 = self.attention2(torch.cat([x2, f3], dim=1), d3.detach())
+        d2 = self.image_pyramid.reconstruct(d3.detach(), p2) #4
+
+        x1 = self.res(x1, (H // 2, W // 2))
+        f2 = self.res(f2, (H // 2, W // 2))
+        f1, p1 = self.attention1(torch.cat([x1, f2], dim=1), d2.detach(), p2.detach()) #2
+        d1 = self.image_pyramid.reconstruct(d2.detach(), p1) #2
+        
+        f1 = self.res(f1, (H, W))
+        _, p0 = self.attention0(f1, d1.detach(), p1.detach()) #2
+        d0 = self.image_pyramid.reconstruct(d1.detach(), p0) #2
+        
+        out = dict()
+        out['saliency'] = [d3, d2, d1, d0]
+        out['laplacian'] = [p2, p1, p0]
+        
+        return out
+    
+    def forward_inference(self, img, img_lr=None):
+        B, _, H, W = img.shape
+        
+        if self.threshold is None:
+            out = self.forward_inspyre(img)
+            d3, d2, d1, d0 = out['saliency']
+            p2, p1, p0     = out['laplacian']
+            
+        elif (H <= self.threshold or W <= self.threshold):
+            if img_lr is not None:
+                out = self.forward_inspyre(img_lr)
+            else:
+                out = self.forward_inspyre(img)
+            d3, d2, d1, d0 = out['saliency']
+            p2, p1, p0     = out['laplacian']
+        
+        else:
+            # LR Saliency Pyramid
+            lr_out = self.forward_inspyre(img_lr)
+            lr_d3, lr_d2, lr_d1, lr_d0 = lr_out['saliency']
+            lr_p2, lr_p1, lr_p0      = lr_out['laplacian']
+                
+            # HR Saliency Pyramid
+            hr_out = self.forward_inspyre(img)
+            hr_d3, hr_d2, hr_d1, hr_d0 = hr_out['saliency']
+            hr_p2, hr_p1, hr_p0      = hr_out['laplacian']
+            
+            # Pyramid Blending
+            d3 = self.ret(lr_d0, hr_d3) 
+            
+            t2 = self.ret(self.transition2(d3), hr_p2)
+            p2 = t2 * hr_p2
+            d2 = self.image_pyramid.reconstruct(d3, p2)
+            
+            t1 = self.ret(self.transition1(d2), hr_p1)
+            p1 = t1 * hr_p1
+            d1 = self.image_pyramid.reconstruct(d2, p1)
+            
+            t0 = self.ret(self.transition0(d1), hr_p0)
+            p0 = t0 * hr_p0
+            d0 = self.image_pyramid.reconstruct(d1, p0)
+            
+        pred = torch.sigmoid(d0)
+        pred = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)
+
+        return pred
+    
+def InSPyReNet_SwinB(depth, pretrained, base_size, **kwargs):
+    return InSPyReNet(SwinB(pretrained=pretrained), [128, 128, 256, 512, 1024], depth, base_size, **kwargs)

transparent_background/Remover.py

@@ -0,0 +1,419 @@
+import os
+import sys
+import tqdm
+import wget
+import gdown
+import torch
+import shutil
+import base64
+import warnings
+import importlib
+
+import numpy as np
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+import albumentations as A
+import albumentations.pytorch as AP
+
+from PIL import Image
+from io import BytesIO
+from packaging import version
+
+filepath = os.path.abspath(__file__)
+repopath = os.path.split(filepath)[0]
+sys.path.append(repopath)
+
+from transparent_background.InSPyReNet import InSPyReNet_SwinB
+from transparent_background.utils import *
+
+class Remover:
+    def __init__(self, mode="base", jit=False, device=None, ckpt=None, resize='static'):
+        """
+        Args:
+            mode   (str): Choose among below options
+                                   base -> slow & large gpu memory required, high quality results
+                                   fast -> resize input into small size for fast computation
+                                   base-nightly -> nightly release for base mode
+            jit    (bool): use TorchScript for fast computation
+            device (str, optional): specifying device for computation. find available GPU resource if not specified.
+            ckpt   (str, optional): specifying model checkpoint. find downloaded checkpoint or try download if not specified.
+            fast   (bool, optional, DEPRECATED): replaced by mode argument. use fast mode if True.
+        """
+        cfg_path = os.environ.get('TRANSPARENT_BACKGROUND_FILE_PATH', os.path.abspath(os.path.expanduser('~')))
+        home_dir = os.path.join(cfg_path, ".transparent-background")
+        os.makedirs(home_dir, exist_ok=True)
+
+        if not os.path.isfile(os.path.join(home_dir, "config.yaml")):
+            shutil.copy(os.path.join(repopath, "config.yaml"), os.path.join(home_dir, "config.yaml"))
+        self.meta = load_config(os.path.join(home_dir, "config.yaml"))[mode]
+
+        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 = home_dir
+            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()
+            ):
+                if self.meta.md5 is not None:
+                    download = True
+
+            if download:
+                if 'drive.google.com' in self.meta.url:
+                    gdown.download(self.meta.url, os.path.join(ckpt_dir, ckpt_name), fuzzy=True, proxy=self.meta.http_proxy)
+                elif 'github.com' in self.meta.url:
+                    wget.download(self.meta.url, os.path.join(ckpt_dir, ckpt_name))
+                else:
+                    raise NotImplementedError('Please use valid URL')
+        else:
+            ckpt_dir, ckpt_name = os.path.split(os.path.abspath(ckpt))
+
+        self.model = InSPyReNet_SwinB(depth=64, pretrained=False, threshold=None, **self.meta)
+        self.model.eval()
+        self.model.load_state_dict(
+            torch.load(os.path.join(ckpt_dir, ckpt_name), map_location="cpu", weights_only=True),
+            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))
+            if resize != 'static':
+                warnings.warn('Resizing method for TorchScript mode only supports static resize. Fallback to static.')
+                resize = 'static'
+
+        resize_tf = None
+        resize_fn = None
+        if resize == 'static':
+            resize_tf = static_resize(self.meta.base_size)
+            resize_fn = A.Resize(*self.meta.base_size)
+        elif resize == 'dynamic':
+            if 'base' not in mode:
+                warnings.warn('Dynamic resizing only supports base and base-nightly mode. It will cause severe performance degradation with fast mode.')
+            resize_tf = dynamic_resize(L=1280)
+            resize_fn = dynamic_resize_a(L=1280)
+        else:
+            raise AttributeError(f'Unsupported resizing method {resize}')
+
+        self.transform = transforms.Compose(
+            [
+                resize_tf,
+                tonumpy(),
+                normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+                totensor(),
+            ]
+        )
+
+        self.cv2_transform = A.Compose(
+            [
+                resize_fn,
+                A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+                AP.ToTensorV2(),
+            ]
+        )
+
+        self.background = {'img': None, 'name': None, 'shape': None}
+        desc = "Mode={}, Device={}, Torchscript={}".format(
+            mode, self.device, "enabled" if jit else "disabled"
+        )
+        print("Settings -> {}".format(desc))
+
+    def process(self, img, type="rgba", threshold=None, reverse=False):
+        """
+        Args:
+            img (PIL.Image or np.ndarray): input image as PIL.Image or np.ndarray type
+            type (str): output type option as below.
+                        'rgba' will generate RGBA output regarding saliency score as an alpha map. 
+                        'green' will change the background with green screen.
+                        'white' will change the background with white color.
+                        '[255, 0, 0]' will change the background with color code [255, 0, 0]. 
+                        'blur' will blur the background.
+                        'overlay' will cover the salient object with translucent green color, and highlight the edges.
+                        Another image file (e.g., 'samples/backgroud.png') will be used as a background, and the object will be overlapped on it.
+            threshold (float or str, optional): produce hard prediction w.r.t specified threshold value (0.0 ~ 1.0)
+        Returns:
+            PIL.Image: output image
+
+        """
+
+        if isinstance(img, np.ndarray):
+            is_numpy = True
+            shape = img.shape[:2]
+            x = self.cv2_transform(image=img)["image"]
+        else:
+            is_numpy = False
+            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()
+
+        if threshold is not None:
+            pred = (pred > float(threshold)).astype(np.float64)
+        if reverse:
+            pred = 1 - pred
+
+        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":
+            if threshold is None:
+                # pymatting is imported here to avoid the overhead in other cases.
+                try:
+                    from pymatting.foreground.estimate_foreground_ml_cupy import estimate_foreground_ml_cupy as estimate_foreground_ml
+                except ImportError:
+                    try:
+                        from pymatting.foreground.estimate_foreground_ml_pyopencl import estimate_foreground_ml_pyopencl as estimate_foreground_ml
+                    except ImportError:
+                        from pymatting import estimate_foreground_ml
+                img = estimate_foreground_ml(img / 255.0, pred)
+                img = 255 * np.clip(img, 0., 1.) + 0.5
+                img = img.astype(np.uint8)
+
+            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:
+            print(type)
+            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 > 0.5) * 255).astype(np.uint8), 50, 100)
+            img[border != 0] = [120, 255, 155]
+
+        elif type.lower().endswith(IMG_EXTS):
+            if self.background['name'] != type:
+                background_img = cv2.cvtColor(cv2.imread(type), cv2.COLOR_BGR2RGB)
+                background_img = cv2.resize(background_img, img.shape[:2][::-1])
+                
+                self.background['img'] = background_img
+                self.background['shape'] = img.shape[:2][::-1]
+                self.background['name'] = type
+            
+            elif self.background['shape'] != img.shape[:2][::-1]:
+                self.background['img'] = cv2.resize(self.background['img'], img.shape[:2][::-1])
+                self.background['shape'] = img.shape[:2][::-1]
+
+            img = img * pred[..., np.newaxis] + self.background['img'] * (
+                1 - pred[..., np.newaxis]
+            )
+
+        if is_numpy:
+            return img.astype(np.uint8)
+        else:
+            return Image.fromarray(img.astype(np.uint8))
+
+def to_base64(image):
+    buffered = BytesIO()
+    image.save(buffered, format="JPEG")
+    base64_img = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return base64_img
+
+def entry_point(out_type, mode, device, ckpt, source, dest, jit, threshold, resize, save_format=None, reverse=False, flet_progress=None, flet_page=None, preview=None, preview_out=None, options=None):
+    warnings.filterwarnings("ignore")
+
+    remover = Remover(mode=mode, jit=jit, device=device, ckpt=ckpt, resize=resize)
+
+    if source.isnumeric() is True:
+        save_dir = None
+        _format = "Webcam"
+        if importlib.util.find_spec('pyvirtualcam') is not None:
+            try:
+                import pyvirtualcam
+                vcam = pyvirtualcam.Camera(width=640, height=480, fps=30)
+            except:
+                vcam = None
+        else:
+            raise ImportError("pyvirtualcam not found. Install with \"pip install transparent-background[webcam]\"")
+
+    elif os.path.isdir(source):
+        save_dir = os.path.join(os.getcwd(), source.split(os.sep)[-1])
+        _format = get_format(os.listdir(source))
+
+    elif os.path.isfile(source):
+        save_dir = os.getcwd()
+        _format = get_format([source])
+
+    else:
+        raise FileNotFoundError("File or directory {} is invalid.".format(source))
+
+    if out_type == "rgba" and _format == "Video":
+        raise AttributeError("type 'rgba' cannot be applied to video input.")
+
+    if dest is not None:
+        save_dir = dest
+
+    if save_dir is not None:
+        os.makedirs(save_dir, exist_ok=True)
+
+    loader = eval(_format + "Loader")(source)
+    frame_progress = tqdm.tqdm(
+        total=len(loader),
+        position=1 if (_format == "Video" and len(loader) > 1) else 0,
+        leave=False,
+        bar_format="{desc:<15}{percentage:3.0f}%|{bar:50}{r_bar}",
+    )
+    sample_progress = (
+        tqdm.tqdm(
+            total=len(loader),
+            desc="Total:",
+            position=0,
+            bar_format="{desc:<15}{percentage:3.0f}%|{bar:50}{r_bar}",
+        )
+        if (_format == "Video" and len(loader) > 1)
+        else None
+    )
+    if flet_progress is not None:
+        assert flet_page is not None
+        flet_progress.value = 0
+        flet_step = 1 / frame_progress.total
+
+    writer = None
+
+    for img, name in loader:
+        filename, ext = os.path.splitext(name)
+        ext = ext[1:]
+        ext = save_format if save_format is not None else ext
+        frame_progress.set_description("{}".format(name))
+        if out_type.lower().endswith(IMG_EXTS):
+            outname = "{}_{}".format(
+                filename,
+                os.path.splitext(os.path.split(out_type)[-1])[0],
+            )
+        else:
+            outname = "{}_{}".format(filename, out_type)
+
+        if reverse:
+            outname += '_reverse'
+
+        if _format == "Video" and writer is None:
+            writer = cv2.VideoWriter(
+                os.path.join(save_dir, f"{outname}.{ext}"),
+                cv2.VideoWriter_fourcc(*"mp4v"),
+                loader.fps,
+                img.size,
+            )
+            writer.set(cv2.VIDEOWRITER_PROP_QUALITY, 100)
+            frame_progress.refresh()
+            frame_progress.reset()
+            frame_progress.total = int(loader.cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            if sample_progress is not None:
+                sample_progress.update()
+
+            if flet_progress is not None:
+                assert flet_page is not None
+                flet_progress.value = 0
+                flet_step = 1 / frame_progress.total
+                flet_progress.update()
+
+        if _format == "Video" and img is None:
+            if writer is not None:
+                writer.release()
+            writer = None
+            continue
+
+        out = remover.process(img, type=out_type, threshold=threshold, reverse=reverse)
+
+        if _format == "Image":
+            if out_type == "rgba" and ext.lower() != 'png':
+                warnings.warn('Output format for rgba mode only supports png format. Fallback to png output.')
+                ext = 'png'
+            out.save(os.path.join(save_dir, f"{outname}.{ext}"))
+        elif _format == "Video" and writer is not None:
+            writer.write(cv2.cvtColor(np.array(out), cv2.COLOR_BGR2RGB))
+        elif _format == "Webcam":
+            if vcam is not None:
+                vcam.send(np.array(out))
+                vcam.sleep_until_next_frame()
+            else:
+                cv2.imshow(
+                    "transparent-background", cv2.cvtColor(np.array(out), cv2.COLOR_BGR2RGB)
+                )
+        frame_progress.update()
+        if flet_progress is not None:
+            flet_progress.value += flet_step
+            flet_progress.update()
+
+            if out_type == 'rgba':
+                o = np.array(out).astype(np.float64)
+                o[:, :, :3] *= (o[:, :, -1:] / 255)
+                out = Image.fromarray(o[:, :, :3].astype(np.uint8))
+
+            preview.src_base64 = to_base64(img.resize((480, 300)).convert('RGB'))
+            preview_out.src_base64 = to_base64(out.resize((480, 300)).convert('RGB'))
+            preview.update()
+            preview_out.update()
+
+        if options is not None and options['abort']:
+            break
+        
+    print("\nDone. Results are saved in {}".format(os.path.abspath(save_dir)))
+
+def console():
+    args = parse_args()
+    entry_point(args.type, args.mode, args.device, args.ckpt, args.source, args.dest, args.jit, args.threshold, args.resize, args.format, args.reverse)

transparent_background/__init__.py

@@ -0,0 +1,2 @@
+from transparent_background.Remover import Remover, console
+from transparent_background.gui import gui

transparent_background/backbones/SwinTransformer.py

@@ -0,0 +1,652 @@
+# --------------------------------------------------------
+# Swin Transformer
+# Copyright (c) 2021 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ze Liu, Yutong Lin, Yixuan Wei
+# --------------------------------------------------------
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+import numpy as np
+from timm.layers import DropPath, to_2tuple, trunc_normal_
+
+class Mlp(nn.Module):
+    """ Multilayer perceptron."""
+
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+def window_partition(x, window_size):
+    """
+    Args:
+        x: (B, H, W, C)
+        window_size (int): window size
+
+    Returns:
+        windows: (num_windows*B, window_size, window_size, C)
+    """
+    B, H, W, C = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    return windows
+
+
+def window_reverse(windows, window_size, H, W):
+    """
+    Args:
+        windows: (num_windows*B, window_size, window_size, C)
+        window_size (int): Window size
+        H (int): Height of image
+        W (int): Width of image
+
+    Returns:
+        x: (B, H, W, C)
+    """
+    B = int(windows.shape[0] / (H * W / window_size / window_size))
+    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return x
+
+
+class WindowAttention(nn.Module):
+    """ Window based multi-head self attention (W-MSA) module with relative position bias.
+    It supports both of shifted and non-shifted window.
+
+    Args:
+        dim (int): Number of input channels.
+        window_size (tuple[int]): The height and width of the window.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
+
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(self.window_size[0])
+        coords_w = torch.arange(self.window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        self.register_buffer("relative_position_index", relative_position_index)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        trunc_normal_(self.relative_position_bias_table, std=.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask=None):
+        """ Forward function.
+
+        Args:
+            x: input features with shape of (num_windows*B, N, C)
+            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
+        """
+        B_, N, C = x.shape
+        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class SwinTransformerBlock(nn.Module):
+    """ Swin Transformer Block.
+
+    Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, dim, num_heads, window_size=7, shift_size=0,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
+
+        self.norm1 = norm_layer(dim)
+        self.attn = WindowAttention(
+            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        self.H = None
+        self.W = None
+
+    def forward(self, x, mask_matrix):
+        """ Forward function.
+
+        Args:
+            x: Input feature, tensor size (B, H*W, C).
+            H, W: Spatial resolution of the input feature.
+            mask_matrix: Attention mask for cyclic shift.
+        """
+        B, L, C = x.shape
+        H, W = self.H, self.W
+        assert L == H * W, "input feature has wrong size"
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(B, H, W, C)
+
+        # pad feature maps to multiples of window size
+        pad_l = pad_t = 0
+        pad_r = (self.window_size - W % self.window_size) % self.window_size
+        pad_b = (self.window_size - H % self.window_size) % self.window_size
+        x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
+        _, Hp, Wp, _ = x.shape
+
+        # cyclic shift
+        if self.shift_size > 0:
+            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+            attn_mask = mask_matrix
+        else:
+            shifted_x = x
+            attn_mask = None
+
+        # partition windows
+        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
+        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
+
+        # W-MSA/SW-MSA
+        attn_windows = self.attn(x_windows, mask=attn_mask)  # nW*B, window_size*window_size, C
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
+        shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp)  # B H' W' C
+
+        # reverse cyclic shift
+        if self.shift_size > 0:
+            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+        else:
+            x = shifted_x
+
+        if pad_r > 0 or pad_b > 0:
+            x = x[:, :H, :W, :].contiguous()
+
+        x = x.view(B, H * W, C)
+
+        # FFN
+        x = shortcut + self.drop_path(x)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x
+
+
+class PatchMerging(nn.Module):
+    """ Patch Merging Layer
+
+    Args:
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+    def __init__(self, dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = norm_layer(4 * dim)
+
+    def forward(self, x, H, W):
+        """ Forward function.
+
+        Args:
+            x: Input feature, tensor size (B, H*W, C).
+            H, W: Spatial resolution of the input feature.
+        """
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+
+        x = x.view(B, H, W, C)
+
+        # padding
+        pad_input = (H % 2 == 1) or (W % 2 == 1)
+        if pad_input:
+            x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))
+
+        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
+        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
+        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
+        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
+        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
+        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
+
+        x = self.norm(x)
+        x = self.reduction(x)
+
+        return x
+
+
+class BasicLayer(nn.Module):
+    """ A basic Swin Transformer layer for one stage.
+
+    Args:
+        dim (int): Number of feature channels
+        depth (int): Depths of this stage.
+        num_heads (int): Number of attention head.
+        window_size (int): Local window size. Default: 7.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+    """
+
+    def __init__(self,
+                 dim,
+                 depth,
+                 num_heads,
+                 window_size=7,
+                 mlp_ratio=4.,
+                 qkv_bias=True,
+                 qk_scale=None,
+                 drop=0.,
+                 attn_drop=0.,
+                 drop_path=0.,
+                 norm_layer=nn.LayerNorm,
+                 downsample=None,
+                 use_checkpoint=False):
+        super().__init__()
+        self.window_size = window_size
+        self.shift_size = window_size // 2
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList([
+            SwinTransformerBlock(
+                dim=dim,
+                num_heads=num_heads,
+                window_size=window_size,
+                shift_size=0 if (i % 2 == 0) else window_size // 2,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop,
+                attn_drop=attn_drop,
+                drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                norm_layer=norm_layer)
+            for i in range(depth)])
+
+        # patch merging layer
+        if downsample is not None:
+            self.downsample = downsample(dim=dim, norm_layer=norm_layer)
+        else:
+            self.downsample = None
+
+    def forward(self, x, H, W):
+        """ Forward function.
+
+        Args:
+            x: Input feature, tensor size (B, H*W, C).
+            H, W: Spatial resolution of the input feature.
+        """
+
+        # calculate attention mask for SW-MSA
+        Hp = int(np.ceil(H / self.window_size)) * self.window_size
+        Wp = int(np.ceil(W / self.window_size)) * self.window_size
+        img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device)  # 1 Hp Wp 1
+        h_slices = (slice(0, -self.window_size),
+                    slice(-self.window_size, -self.shift_size),
+                    slice(-self.shift_size, None))
+        w_slices = (slice(0, -self.window_size),
+                    slice(-self.window_size, -self.shift_size),
+                    slice(-self.shift_size, None))
+        cnt = 0
+        for h in h_slices:
+            for w in w_slices:
+                img_mask[:, h, w, :] = cnt
+                cnt += 1
+
+        mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
+        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+
+        for blk in self.blocks:
+            blk.H, blk.W = H, W
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x, attn_mask)
+            else:
+                x = blk(x, attn_mask)
+        if self.downsample is not None:
+            x_down = self.downsample(x, H, W)
+            Wh, Ww = (H + 1) // 2, (W + 1) // 2
+            return x, H, W, x_down, Wh, Ww
+        else:
+            return x, H, W, x, H, W
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+
+    Args:
+        patch_size (int): Patch token size. Default: 4.
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+    """
+
+    def __init__(self, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        patch_size = to_2tuple(patch_size)
+        self.patch_size = patch_size
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        """Forward function."""
+        # padding
+        _, _, H, W = x.size()
+        if W % self.patch_size[1] != 0:
+            x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))
+        if H % self.patch_size[0] != 0:
+            x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))
+
+        x = self.proj(x)  # B C Wh Ww
+        if self.norm is not None:
+            Wh, Ww = x.size(2), x.size(3)
+            x = x.flatten(2).transpose(1, 2)
+            x = self.norm(x)
+            x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)
+
+        return x
+
+
+class SwinTransformer(nn.Module):
+    """ Swin Transformer backbone.
+        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -
+          https://arxiv.org/pdf/2103.14030
+
+    Args:
+        pretrain_img_size (int): Input image size for training the pretrained model,
+            used in absolute postion embedding. Default 224.
+        patch_size (int | tuple(int)): Patch size. Default: 4.
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        depths (tuple[int]): Depths of each Swin Transformer stage.
+        num_heads (tuple[int]): Number of attention head of each stage.
+        window_size (int): Window size. Default: 7.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set.
+        drop_rate (float): Dropout rate.
+        attn_drop_rate (float): Attention dropout rate. Default: 0.
+        drop_path_rate (float): Stochastic depth rate. Default: 0.2.
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False.
+        patch_norm (bool): If True, add normalization after patch embedding. Default: True.
+        out_indices (Sequence[int]): Output from which stages.
+        frozen_stages (int): Stages to be frozen (stop grad and set eval mode).
+            -1 means not freezing any parameters.
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+    """
+
+    def __init__(self,
+                 pretrain_img_size=224,
+                 patch_size=4,
+                 in_chans=3,
+                 embed_dim=96,
+                 depths=[2, 2, 6, 2],
+                 num_heads=[3, 6, 12, 24],
+                 window_size=7,
+                 mlp_ratio=4.,
+                 qkv_bias=True,
+                 qk_scale=None,
+                 drop_rate=0.,
+                 attn_drop_rate=0.,
+                 drop_path_rate=0.2,
+                 norm_layer=nn.LayerNorm,
+                 ape=False,
+                 patch_norm=True,
+                 out_indices=(0, 1, 2, 3),
+                 frozen_stages=-1,
+                 use_checkpoint=False):
+        super().__init__()
+
+        self.pretrain_img_size = pretrain_img_size
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.ape = ape
+        self.patch_norm = patch_norm
+        self.out_indices = out_indices
+        self.frozen_stages = frozen_stages
+
+        # split image into non-overlapping patches
+        self.patch_embed = PatchEmbed(
+            patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None)
+
+        # absolute position embedding
+        if self.ape:
+            pretrain_img_size = to_2tuple(pretrain_img_size)
+            patch_size = to_2tuple(patch_size)
+            patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]]
+
+            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1]))
+            trunc_normal_(self.absolute_pos_embed, std=.02)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
+
+        # build layers
+        self.layers = nn.ModuleList()
+        for i_layer in range(self.num_layers):
+            layer = BasicLayer(
+                dim=int(embed_dim * 2 ** i_layer),
+                depth=depths[i_layer],
+                num_heads=num_heads[i_layer],
+                window_size=window_size,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop_rate,
+                attn_drop=attn_drop_rate,
+                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
+                norm_layer=norm_layer,
+                downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
+                use_checkpoint=use_checkpoint)
+            self.layers.append(layer)
+
+        num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]
+        self.num_features = num_features
+
+        # add a norm layer for each output
+        for i_layer in out_indices:
+            layer = norm_layer(num_features[i_layer])
+            layer_name = f'norm{i_layer}'
+            self.add_module(layer_name, layer)
+
+        self._freeze_stages()
+
+    def _freeze_stages(self):
+        if self.frozen_stages >= 0:
+            self.patch_embed.eval()
+            for param in self.patch_embed.parameters():
+                param.requires_grad = False
+
+        if self.frozen_stages >= 1 and self.ape:
+            self.absolute_pos_embed.requires_grad = False
+
+        if self.frozen_stages >= 2:
+            self.pos_drop.eval()
+            for i in range(0, self.frozen_stages - 1):
+                m = self.layers[i]
+                m.eval()
+                for param in m.parameters():
+                    param.requires_grad = False
+
+    def init_weights(self, pretrained=None):
+        """Initialize the weights in backbone.
+
+        Args:
+            pretrained (str, optional): Path to pre-trained weights.
+                Defaults to None.
+        """
+
+        def _init_weights(m):
+            if isinstance(m, nn.Linear):
+                trunc_normal_(m.weight, std=.02)
+                if isinstance(m, nn.Linear) and m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.LayerNorm):
+                nn.init.constant_(m.bias, 0)
+                nn.init.constant_(m.weight, 1.0)
+
+        if isinstance(pretrained, str):
+            self.apply(_init_weights)
+            logger = get_root_logger()
+            load_checkpoint(self, pretrained, strict=False, logger=logger)
+        elif pretrained is None:
+            self.apply(_init_weights)
+        else:
+            raise TypeError('pretrained must be a str or None')
+
+    def forward(self, x):
+        """Forward function."""
+        x = self.patch_embed(x)
+
+        Wh, Ww = x.size(2), x.size(3)
+        if self.ape:
+            # interpolate the position embedding to the corresponding size
+            absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wh, Ww), mode='bicubic')
+            x = (x + absolute_pos_embed) # B Wh*Ww C
+            
+        outs = [x.contiguous()]
+        x = x.flatten(2).transpose(1, 2)
+        x = self.pos_drop(x)
+        for i in range(self.num_layers):
+            layer = self.layers[i]
+            x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)
+
+            if i in self.out_indices:
+                norm_layer = getattr(self, f'norm{i}')
+                x_out = norm_layer(x_out)
+
+                out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
+                outs.append(out)
+
+        return tuple(outs)
+
+    def train(self, mode=True):
+        """Convert the model into training mode while keep layers freezed."""
+        super(SwinTransformer, self).train(mode)
+        self._freeze_stages()
+
+def SwinT(pretrained=True):
+    model = SwinTransformer(embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], window_size=7)
+    if pretrained is True:
+        model.load_state_dict(torch.load('data/backbone_ckpt/swin_tiny_patch4_window7_224.pth', map_location='cpu')['model'], strict=False)
+        
+    return model
+
+def SwinS(pretrained=True):
+    model = SwinTransformer(embed_dim=96, depths=[2, 2, 18, 2], num_heads=[3, 6, 12, 24], window_size=7)
+    if pretrained is True:
+        model.load_state_dict(torch.load('data/backbone_ckpt/swin_small_patch4_window7_224.pth', map_location='cpu')['model'], strict=False)
+        
+    return model
+
+def SwinB(pretrained=True):
+    model = SwinTransformer(embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12)
+    if pretrained is True:
+        model.load_state_dict(torch.load('data/backbone_ckpt/swin_base_patch4_window12_384_22kto1k.pth', map_location='cpu')['model'], strict=False)
+        
+    return model
+
+def SwinL(pretrained=True):
+    model = SwinTransformer(embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=12)
+    if pretrained is True:
+        model.load_state_dict(torch.load('data/backbone_ckpt/swin_large_patch4_window12_384_22kto1k.pth', map_location='cpu')['model'], strict=False)
+
+    return model
+

transparent_background/config.yaml

@@ -0,0 +1,21 @@
+base:
+  url: "https://github.com/plemeri/transparent-background/releases/download/1.2.12/ckpt_base.pth"
+  md5: "d692e3dd5fa1b9658949d452bebf1cda"
+  ckpt_name: "ckpt_base.pth"
+  http_proxy: NULL
+  base_size: [1024, 1024]
+
+
+fast:
+  url: "https://github.com/plemeri/transparent-background/releases/download/1.2.12/ckpt_fast.pth"
+  md5: "9efdbfbcc49b79ef0f7891c83d2fd52f"
+  ckpt_name: "ckpt_fast.pth"
+  http_proxy: NULL
+  base_size: [384, 384]
+
+base-nightly:
+  url: "https://github.com/plemeri/transparent-background/releases/download/1.2.12/ckpt_base_nightly.pth"
+  md5: NULL
+  ckpt_name: "ckpt_base_nightly.pth"
+  http_proxy: NULL
+  base_size: [1024, 1024]

transparent_background/gui.py

@@ -0,0 +1,344 @@
+import flet as ft
+from flet import (
+    ElevatedButton,
+    FilePicker,
+    FilePickerResultEvent,
+    Page,
+    Row,
+    Text,
+    icons,
+)
+import os
+import torch
+import logging
+
+from transparent_background.utils import *
+from transparent_background.Remover import *
+
+logging.basicConfig(level=logging.WARN)
+logging.getLogger("flet_runtime").setLevel(logging.WARN)
+
+options = {
+    'output_type':'rgba',
+    'mode':'base',
+    'device':get_backend(),
+    'r' : 0, 
+    'g' : 0, 
+    'b' : 0,
+    'color' : "[0, 0, 0]",
+    'ckpt':None,
+    'threshold':None,
+    'reverse': False,
+    'resize': 'static',
+    'format': None,
+    'source':None,
+    'dest':None,
+    'use_custom':False,
+    'jit':False,
+    'abort':False,
+}
+
+def is_float(str):
+    if str is None:
+        return False
+    try:
+        tmp = float(str)
+        return True
+    except ValueError:
+        return False
+
+def main(page):
+    def theme_changed(e):
+        page.theme_mode = (
+            ft.ThemeMode.DARK
+            if page.theme_mode == ft.ThemeMode.LIGHT
+            else ft.ThemeMode.LIGHT
+        )
+        page.update()
+
+    def checkbox_changed(e):
+        options['jit'] = jit_check.value
+        options['reverse'] = reverse_check.value
+        page.update()
+
+    def dropdown_changed(e):
+        options['output_type'] = type_dropdown.value
+        options['mode'] = mode_dropdown.value
+        options['device'] = device_dropdown.value
+        options['resize'] = resize_dropdown.value
+        # options['format'] = format_dropdown.value
+
+        if options['output_type'] == 'custom' and not options['use_custom']:
+            page.insert(1, ft.Row([r_field, g_field, b_field]))
+            options['use_custom']=True
+
+        elif options['output_type'] != 'custom' and options['use_custom']:
+            options['use_custom']=False
+            page.remove_at(1)
+
+        output_text.value = 'Type: {}, Mode: {}, Device: {}, Threshold: {}, Resize: {}, Format: {}'.format(options['output_type'], options['mode'], options['device'], options['threshold'], options['resize'], options['format'])
+        page.update()
+
+    def color_changed(e):
+        options['r'] = int(r_field.value) if len(r_field.value) > 0 and r_field.value.isdigit() else 0
+        options['g'] = int(g_field.value) if len(g_field.value) > 0 and g_field.value.isdigit() else 0
+        options['b'] = int(b_field.value) if len(b_field.value) > 0 and b_field.value.isdigit() else 0
+        options['color'] = str([options['r'], options['g'], options['b']])
+        output_text.value = 'Type: {}, Mode: {}, Device: {}, Threshold: {}, Resize: {}, Format: {}'.format(options['output_type'], options['color'], options['device'], options['threshold'], options['resize'], options['format'])
+        page.update()
+
+    def threshold_changed(e):
+        options['threshold'] = float(threshold_field.value) if len(threshold_field.value) > 0 and is_float(threshold_field.value) else None
+        options['threshold'] = None if is_float(options['threshold']) and (options['threshold'] < 0 or options['threshold'] > 1) else options['threshold']
+        output_text.value = 'Type: {}, Mode: {}, Device: {}, Threshold: {}, Resize: {}, Format: {}'.format(options['output_type'], options['mode'], options['device'], options['threshold'], options['resize'], options['format'])
+        page.update()
+    
+    def format_changed(e):
+        options['format'] = format_field.value if format_field.value.endswith(IMG_EXTS) or format_field.value.endswith(VID_EXTS) else None
+        output_text.value = 'Type: {}, Mode: {}, Device: {}, Threshold: {}, Resize: {}, Format: {}'.format(options['output_type'], options['mode'], options['device'], options['threshold'], options['resize'], options['format'])
+        page.update()
+
+    def pick_files_result(e: FilePickerResultEvent):
+        file_path.update()
+        options['source'] = e.files[0].path if e.files else 'Not Selected'
+        file_path.value = options['source']
+        file_path.update()
+        if options['dest'] is None:
+            options['dest'] = os.path.split(options['source'])[0]
+            dest_path.value = options['dest']
+            dest_path.update()
+
+    # Open directory dialog
+    def get_directory_result(e: FilePickerResultEvent):
+        options['source'] = e.path if e.path else 'Not Selected'
+        file_path.value = options['source']
+        file_path.update()
+        if options['dest'] is None:
+            options['dest'] = os.path.split(options['source'])[0]
+            dest_path.value = options['dest']
+            dest_path.update()
+
+    def get_dest_result(e: FilePickerResultEvent):
+        options['dest'] = e.path if e.path else 'Not Selected'
+        dest_path.value = options['dest']
+        dest_path.update()
+
+    def process(e):
+        output_type = options['output_type']
+        output_type = options['color'] if output_type == 'custom' else output_type
+        options['abort'] = False
+        entry_point(output_type, options['mode'], options['device'], options['ckpt'], options['source'], options['dest'], options['jit'], options['threshold'], options['resize'], options['format'], options['reverse'], progress_ring, page, preview, preview_out, options)
+
+    def click_abort(e):
+        options['abort'] = True
+        page.update()
+
+    page.window_width = 1000
+    page.window_height = 650
+    page.window_resizable = False
+
+    page.theme_mode = ft.ThemeMode.LIGHT
+    c = ft.Switch(label="Dark mode", on_change=theme_changed)
+
+    output_text = ft.Text(color=ft.colors.BLACK)
+    output_text.value = 'Type: {}, Mode: {}, Device: {}, Threshold: {}, Resize: {}, Format: {}'.format(options['output_type'], options['mode'], options['device'], options['threshold'], options['resize'], options['format'])
+    output_text_container = ft.Container(
+                    content=output_text,
+                    margin=10,
+                    padding=10,
+                    bgcolor=ft.colors.GREEN_100,
+                    border_radius=10,
+                )
+
+    jit_check = ft.Checkbox(label="use torchscript", value=False, on_change=checkbox_changed)
+    reverse_check = ft.Checkbox(label="reverse", value=False, on_change=checkbox_changed)
+
+    type_dropdown = ft.Dropdown(
+        label='type',
+        width=200,
+        hint_text='output type',
+        on_change=dropdown_changed,
+        options=[
+            ft.dropdown.Option("rgba"),
+            ft.dropdown.Option("map"),
+            ft.dropdown.Option("green"),
+            ft.dropdown.Option("white"),
+            ft.dropdown.Option("blur"),
+            ft.dropdown.Option("overlay"),
+            ft.dropdown.Option("custom"),
+        ],
+    )
+    type_dropdown.value = options['output_type']
+
+    resize_dropdown = ft.Dropdown(
+        label='resize',
+        width=200,
+        hint_text='resize method',
+        on_change=dropdown_changed,
+        options=[
+            ft.dropdown.Option("static"),
+            ft.dropdown.Option("dynamic"),
+        ],
+    )
+    resize_dropdown.value = options['resize']
+
+    Remover() # init once
+
+    cfg_path = os.environ.get('TRANSPARENT_BACKGROUND_FILE_PATH', os.path.abspath(os.path.expanduser('~')))
+    home_dir = os.path.join(cfg_path, ".transparent-background")
+    configs = load_config(os.path.join(home_dir, "config.yaml"))
+
+    mode_dropdown = ft.Dropdown(
+        label='mode',
+        width=150,
+        hint_text='mode',
+        on_change=dropdown_changed,
+        options=[ft.dropdown.Option(key) for key in configs.keys()],
+    )
+    mode_dropdown.value = options['mode']
+
+    device_options = [ft.dropdown.Option("cpu")] 
+    device_options += [ft.dropdown.Option("cuda:{}".format(i)) for i in range(torch.cuda.device_count())]
+    device_options += ['mps:0'] if torch.backends.mps.is_available() else []
+
+    device_dropdown = ft.Dropdown(
+        label='device',
+        width=150,
+        hint_text='device',
+        on_change=dropdown_changed,
+        options=device_options
+    )
+    device_dropdown.value=options['device']
+
+    r_field = ft.TextField(width=60, label='R', on_change=color_changed)
+    g_field = ft.TextField(width=60, label='G', on_change=color_changed)
+    b_field = ft.TextField(width=60, label='B', on_change=color_changed)
+
+    r_field.value=str(options['r'])
+    g_field.value=str(options['g'])
+    b_field.value=str(options['b'])
+
+    threshold_field = ft.TextField(width=150, label='threshold', on_change=threshold_changed)
+    threshold_field.value = None
+
+    format_field = ft.TextField(width=100, label='format', on_change=format_changed)
+    format_field.value = None
+
+    page.add(
+        ft.Row(
+            [
+                ft.Image(src='https://raw.githubusercontent.com/plemeri/transparent-background/main/figures/logo.png', width=100, height=100),
+                ft.Column(
+                    [
+                        ft.Row([c, jit_check, reverse_check, output_text_container]),
+                        ft.Row([type_dropdown, mode_dropdown, device_dropdown, resize_dropdown, threshold_field, format_field])
+                    ]
+                )
+            ]
+        )
+    )
+
+    pick_files_dialog = FilePicker(on_result=pick_files_result)
+
+    get_directory_dialog = FilePicker(on_result=get_directory_result)
+    file_path = Text(color=ft.colors.BLACK)
+    file_path.value = 'Input file or directory will be displayed'
+    file_path_container = ft.Container(
+                content=file_path,
+                margin=10,
+                padding=10,
+                bgcolor=ft.colors.AMBER,
+                border_radius=10,
+            )
+
+    get_dest_dialog = FilePicker(on_result=get_dest_result)
+    dest_path = Text(color=ft.colors.BLACK)
+    dest_path.value = 'Output file or directory will be displayed'
+    dest_path_container = ft.Container(
+            content=dest_path,
+            margin=10,
+            padding=10,
+            bgcolor=ft.colors.CYAN_200,
+            border_radius=10,
+        )
+
+    # hide all dialogs in overlay
+    page.overlay.extend([pick_files_dialog, get_directory_dialog, get_dest_dialog])
+    #progress_ring = ft.ProgressRing(width=16, height=16, stroke_width = 2)
+    progress_ring = ft.ProgressBar(width=200, color='amber', bgcolor='#eeeeee')
+    progress_ring.value = 0
+
+    preview = ft.Image(src=".preview.png", )
+    preview_out = ft.Image(src=".preview_out.png")
+
+    page.add(
+        Row(
+            [
+                ElevatedButton(
+                    "Open File",
+                    icon=icons.UPLOAD_FILE,
+                    on_click=lambda _: pick_files_dialog.pick_files(
+                        allow_multiple=False
+                    ),
+                ),
+                ElevatedButton(
+                    "Open Directory",
+                    icon=icons.FOLDER_OPEN,
+                    on_click=lambda _: get_directory_dialog.get_directory_path(),
+                    disabled=page.web,
+                ),
+                file_path_container,
+            ]
+        ),
+        Row(
+            [
+                ElevatedButton(
+                    "Open Destination",
+                    icon=icons.FOLDER_OPEN,
+                    on_click=lambda _: get_dest_dialog.get_directory_path(),
+                    disabled=page.web,
+                ),
+                dest_path_container
+            ]
+        ),
+        Row(
+            [
+                ElevatedButton(
+                    "Process",
+                    icon=icons.SEND,
+                    on_click=process,
+                    disabled=page.web,
+                ),
+                ElevatedButton(
+                    "Stop",
+                    icon=icons.STOP,
+                    on_click=click_abort,
+                    disabled=page.web,
+                ),
+                progress_ring
+            ]
+        ),
+    )
+
+    page.add(
+        Row(
+            [
+                preview,
+                preview_out
+            ]
+        )
+    )
+
+
+def gui():
+    ft.app(target=main)
+
+    if os.path.isfile('.preview.png'):
+        os.remove('.preview.png')
+
+    if os.path.isfile('.preview_out.png'):
+        os.remove('.preview_out.png')
+
+
+if __name__ == "__main__":
+    gui()

transparent_background/modules/attention_module.py

@@ -0,0 +1,103 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torch.nn.parameter import Parameter
+from operator import xor
+from typing import Optional
+
+from transparent_background.modules.layers import *
+
+class SICA(nn.Module):
+    def __init__(self, in_channel, out_channel=1, depth=64, base_size=None, stage=None, lmap_in=False):
+        super(SICA, self).__init__()
+        self.in_channel = in_channel
+        self.depth = depth
+        self.lmap_in = lmap_in
+        if base_size is not None and stage is not None:
+            self.stage_size = (base_size[0] // (2 ** stage), base_size[1] // (2 ** stage))
+        else:
+            self.stage_size = None
+        
+        self.conv_query = nn.Sequential(Conv2d(in_channel, depth, 3, relu=True),
+                                        Conv2d(depth, depth, 3, relu=True))
+        self.conv_key   = nn.Sequential(Conv2d(in_channel, depth, 1, relu=True),
+                                        Conv2d(depth, depth, 1, relu=True))
+        self.conv_value = nn.Sequential(Conv2d(in_channel, depth, 1, relu=True),
+                                        Conv2d(depth, depth, 1, relu=True))
+
+        if self.lmap_in is True:
+            self.ctx = 5
+        else:
+            self.ctx = 3
+
+        self.conv_out1 = Conv2d(depth, depth, 3, relu=True)
+        self.conv_out2 = Conv2d(in_channel + depth, depth, 3, relu=True)
+        self.conv_out3 = Conv2d(depth, depth, 3, relu=True)
+        self.conv_out4 = Conv2d(depth, out_channel, 1)
+
+        self.threshold = Parameter(torch.tensor([0.5]))
+        
+        if self.lmap_in is True:
+            self.lthreshold = Parameter(torch.tensor([0.5]))
+
+    def forward(self, x, smap, lmap: Optional[torch.Tensor]=None):
+        assert not xor(self.lmap_in is True, lmap is not None)
+        b, c, h, w = x.shape
+        
+        # compute class probability
+        smap = F.interpolate(smap, size=x.shape[-2:], mode='bilinear', align_corners=False)
+        smap = torch.sigmoid(smap)
+        p = smap - self.threshold
+
+        fg = torch.clip(p, 0, 1) # foreground
+        bg = torch.clip(-p, 0, 1) # background
+        cg = self.threshold - torch.abs(p) # confusion area
+
+        if self.lmap_in is True and lmap is not None:
+            lmap = F.interpolate(lmap, size=x.shape[-2:], mode='bilinear', align_corners=False)
+            lmap = torch.sigmoid(lmap)
+            lp = lmap - self.lthreshold
+            fp = torch.clip(lp, 0, 1) # foreground
+            bp = torch.clip(-lp, 0, 1) # background
+
+            prob = [fg, bg, cg, fp, bp]
+        else:
+            prob = [fg, bg, cg]
+
+        prob = torch.cat(prob, dim=1)
+
+        # reshape feature & prob
+        if self.stage_size is not None:
+            shape = self.stage_size
+            shape_mul = self.stage_size[0] * self.stage_size[1]
+        else:
+            shape = (h, w)
+            shape_mul = h * w        
+        
+        f = F.interpolate(x, size=shape, mode='bilinear', align_corners=False).view(b, shape_mul, -1)
+        prob = F.interpolate(prob, size=shape, mode='bilinear', align_corners=False).view(b, self.ctx, shape_mul)
+        
+        # compute context vector
+        context = torch.bmm(prob, f).permute(0, 2, 1).unsqueeze(3) # b, 3, c
+
+        # k q v compute
+        query = self.conv_query(x).view(b, self.depth, -1).permute(0, 2, 1)
+        key = self.conv_key(context).view(b, self.depth, -1)
+        value = self.conv_value(context).view(b, self.depth, -1).permute(0, 2, 1)
+
+        # compute similarity map
+        sim = torch.bmm(query, key) # b, hw, c x b, c, 2
+        sim = (self.depth ** -.5) * sim
+        sim = F.softmax(sim, dim=-1)
+
+        # compute refined feature
+        context = torch.bmm(sim, value).permute(0, 2, 1).contiguous().view(b, -1, h, w)
+        context = self.conv_out1(context)
+        
+        x = torch.cat([x, context], dim=1)
+        x = self.conv_out2(x)
+        x = self.conv_out3(x)
+        out = self.conv_out4(x)
+
+        return x, out

transparent_background/modules/context_module.py

@@ -0,0 +1,54 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .layers import *
+
+class PAA_kernel(nn.Module):
+    def __init__(self, in_channel, out_channel, receptive_size, stage_size=None):
+        super(PAA_kernel, self).__init__()
+        self.conv0 = Conv2d(in_channel, out_channel, 1)
+        self.conv1 = Conv2d(out_channel, out_channel, kernel_size=(1, receptive_size))
+        self.conv2 = Conv2d(out_channel, out_channel, kernel_size=(receptive_size, 1))
+        self.conv3 = Conv2d(out_channel, out_channel, 3, dilation=receptive_size)
+        self.Hattn = SelfAttention(out_channel, 'h', stage_size[0] if stage_size is not None else None)
+        self.Wattn = SelfAttention(out_channel, 'w', stage_size[1] if stage_size is not None else None)
+
+    def forward(self, x):
+        x = self.conv0(x)
+        x = self.conv1(x)
+        x = self.conv2(x)
+
+        Hx = self.Hattn(x)
+        Wx = self.Wattn(x)
+
+        x = self.conv3(Hx + Wx)
+        return x
+
+class PAA_e(nn.Module):
+    def __init__(self, in_channel, out_channel, base_size=None, stage=None):
+        super(PAA_e, self).__init__()
+        self.relu = nn.ReLU(True)
+        if base_size is not None and stage is not None:
+            self.stage_size = (base_size[0] // (2 ** stage), base_size[1] // (2 ** stage))
+        else:
+            self.stage_size = None
+
+        self.branch0 = Conv2d(in_channel, out_channel, 1)
+        self.branch1 = PAA_kernel(in_channel, out_channel, 3, self.stage_size)
+        self.branch2 = PAA_kernel(in_channel, out_channel, 5, self.stage_size)
+        self.branch3 = PAA_kernel(in_channel, out_channel, 7, self.stage_size)
+
+        self.conv_cat = Conv2d(4 * out_channel, out_channel, 3)
+        self.conv_res = Conv2d(in_channel, out_channel, 1)
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        x2 = self.branch2(x)
+        x3 = self.branch3(x)
+
+        x_cat = self.conv_cat(torch.cat((x0, x1, x2, x3), 1))
+        x = self.relu(x_cat + self.conv_res(x))
+
+        return x

transparent_background/modules/decoder_module.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .layers import *
+class PAA_d(nn.Module):
+    def __init__(self, in_channel, out_channel=1, depth=64, base_size=None, stage=None):
+        super(PAA_d, self).__init__()
+        self.conv1 = Conv2d(in_channel ,depth, 3)
+        self.conv2 = Conv2d(depth, depth, 3)
+        self.conv3 = Conv2d(depth, depth, 3)
+        self.conv4 = Conv2d(depth, depth, 3)
+        self.conv5 = Conv2d(depth, out_channel, 3, bn=False)
+        
+        self.base_size = base_size
+        self.stage = stage
+        
+        if base_size is not None and stage is not None:
+            self.stage_size = (base_size[0] // (2 ** stage), base_size[1] // (2 ** stage))
+        else:
+            self.stage_size = [None, None]
+
+        self.Hattn = SelfAttention(depth, 'h', self.stage_size[0])
+        self.Wattn = SelfAttention(depth, 'w', self.stage_size[1])
+
+        self.upsample = lambda img, size: F.interpolate(img, size=size, mode='bilinear', align_corners=True)
+
+    def forward(self, fs): #f3 f4 f5 -> f3 f2 f1
+        fx = fs[0]
+        for i in range(1, len(fs)):
+            fs[i] = self.upsample(fs[i], fx.shape[-2:])
+        fx = torch.cat(fs[::-1], dim=1)
+
+        fx = self.conv1(fx)
+
+        Hfx = self.Hattn(fx)
+        Wfx = self.Wattn(fx)
+
+        fx = self.conv2(Hfx + Wfx)
+        fx = self.conv3(fx)
+        fx = self.conv4(fx)
+        out = self.conv5(fx)
+
+        return fx, out

transparent_background/modules/layers.py

@@ -0,0 +1,171 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import cv2
+import numpy as np
+
+from kornia.morphology import dilation, erosion
+from torch.nn.parameter import Parameter
+
+class ImagePyramid:
+    def __init__(self, ksize=7, sigma=1, channels=1):
+        self.ksize = ksize
+        self.sigma = sigma
+        self.channels = channels
+
+        k = cv2.getGaussianKernel(ksize, sigma)
+        k = np.outer(k, k)
+        k = torch.tensor(k).float()
+        self.kernel = k.repeat(channels, 1, 1, 1)
+        
+    def to(self, device):
+        self.kernel = self.kernel.to(device)
+        return self
+        
+    def cuda(self, idx=None):
+        if idx is None:
+            idx = torch.cuda.current_device()
+            
+        self.to(device="cuda:{}".format(idx))
+        return self
+
+    def expand(self, x):
+        z = torch.zeros_like(x)
+        x = torch.cat([x, z, z, z], dim=1)
+        x = F.pixel_shuffle(x, 2)
+        x = F.pad(x, (self.ksize // 2, ) * 4, mode='reflect')
+        x = F.conv2d(x, self.kernel * 4, groups=self.channels)
+        return x
+
+    def reduce(self, x):
+        x = F.pad(x, (self.ksize // 2, ) * 4, mode='reflect')
+        x = F.conv2d(x, self.kernel, groups=self.channels)
+        x = x[:, :, ::2, ::2]
+        return x
+
+    def deconstruct(self, x):
+        reduced_x = self.reduce(x)
+        expanded_reduced_x = self.expand(reduced_x)
+
+        if x.shape != expanded_reduced_x.shape:
+            expanded_reduced_x = F.interpolate(expanded_reduced_x, x.shape[-2:])
+
+        laplacian_x = x - expanded_reduced_x
+        return reduced_x, laplacian_x
+
+    def reconstruct(self, x, laplacian_x):
+        expanded_x = self.expand(x)
+        if laplacian_x.shape != expanded_x:
+            laplacian_x = F.interpolate(laplacian_x, expanded_x.shape[-2:], mode='bilinear', align_corners=True)
+        return expanded_x + laplacian_x
+
+class Transition:
+    def __init__(self, k=3):
+        self.kernel = torch.tensor(cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))).float()
+        
+    def to(self, device):
+        self.kernel = self.kernel.to(device)
+        return self
+        
+    def cuda(self, idx=0):
+        self.to(device="cuda:{}".format(idx))
+        return self
+        
+    def __call__(self, x):
+        x = torch.sigmoid(x)
+        dx = dilation(x, self.kernel)
+        ex = erosion(x, self.kernel)
+        
+        return ((dx - ex) > .5).float()
+
+class Conv2d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, groups=1, padding='same', bias=False, bn=True, relu=False):
+        super(Conv2d, self).__init__()
+        if '__iter__' not in dir(kernel_size):
+            kernel_size = (kernel_size, kernel_size)
+        if '__iter__' not in dir(stride):
+            stride = (stride, stride)
+        if '__iter__' not in dir(dilation):
+            dilation = (dilation, dilation)
+
+        if padding == 'same':
+            width_pad_size = kernel_size[0] + (kernel_size[0] - 1) * (dilation[0] - 1)
+            height_pad_size = kernel_size[1] + (kernel_size[1] - 1) * (dilation[1] - 1)
+        elif padding == 'valid':
+            width_pad_size = 0
+            height_pad_size = 0
+        else:
+            if '__iter__' in dir(padding):
+                width_pad_size = padding[0] * 2
+                height_pad_size = padding[1] * 2
+            else:
+                width_pad_size = padding * 2
+                height_pad_size = padding * 2
+
+        width_pad_size = width_pad_size // 2 + (width_pad_size % 2 - 1)
+        height_pad_size = height_pad_size // 2 + (height_pad_size % 2 - 1)
+        pad_size = (width_pad_size, height_pad_size)
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad_size, dilation, groups, bias=bias)
+        self.reset_parameters()
+
+        if bn is True:
+            self.bn = nn.BatchNorm2d(out_channels)
+        else:
+            self.bn = None
+        
+        if relu is True:
+            self.relu = nn.ReLU(inplace=True)
+        else:
+            self.relu = None
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.bn is not None:
+            x = self.bn(x)
+        if self.relu is not None:
+            x = self.relu(x)
+        return x
+
+    def reset_parameters(self):
+        nn.init.kaiming_normal_(self.conv.weight)
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, in_channels, mode='hw', stage_size=None):
+        super(SelfAttention, self).__init__()
+
+        self.mode = mode
+
+        self.query_conv = Conv2d(in_channels, in_channels // 8, kernel_size=(1, 1))
+        self.key_conv = Conv2d(in_channels, in_channels // 8, kernel_size=(1, 1))
+        self.value_conv = Conv2d(in_channels, in_channels, kernel_size=(1, 1))
+
+        self.gamma = Parameter(torch.zeros(1))
+        self.softmax = nn.Softmax(dim=-1)
+        
+        self.stage_size = stage_size
+
+    def forward(self, x):
+        batch_size, channel, height, width = x.size()
+
+        axis = 1
+        if 'h' in self.mode:
+            axis *= height
+        if 'w' in self.mode:
+            axis *= width
+
+        view = (batch_size, -1, axis)
+
+        projected_query = self.query_conv(x).view(*view).permute(0, 2, 1)
+        projected_key = self.key_conv(x).view(*view)
+
+        attention_map = torch.bmm(projected_query, projected_key)
+        attention = self.softmax(attention_map)
+        projected_value = self.value_conv(x).view(*view)
+
+        out = torch.bmm(projected_value, attention.permute(0, 2, 1))
+        out = out.view(batch_size, channel, height, width)
+
+        out = self.gamma * out + x
+        return out

transparent_background/utils.py

@@ -0,0 +1,261 @@
+import os
+import re
+import cv2
+import yaml
+import torch
+import hashlib
+import argparse
+
+import albumentations as A
+from albumentations.core.transforms_interface import ImageOnlyTransform
+
+import numpy as np
+
+from PIL import Image
+from threading import Thread
+from easydict import EasyDict
+
+VID_EXTS = ('mp4', 'avi', 'h264', 'mkv', 'mov', 'flv', 'wmv', 'webm', 'ts', 'm4v', 'vob', '3gp', '3g2', 'rm', 'rmvb', 'ogv', 'ogg', 'drc', 'gif', 'gifv', 'mng', 'avi', 'mov', 'qt', 'wmv', 'yuv', 'rm', 'rmvb', 'asf', 'amv', 'mp4', 'm4p', 'm4v', 'mpg', 'mp2', 'mpeg', 'mpe', 'mpv', 'mpg', 'mpeg', 'm2v', 'm4v', 'svi', '3gp', '3g2', 'mxf', 'roq', 'nsv', 'flv', 'f4v', 'f4p', 'f4a', 'f4b')
+IMG_EXTS = ('jpg', 'jpeg', 'bmp', 'png', 'ppm', 'pgm', 'pbm', 'pnm', 'webp', 'sr', 'ras', 'tiff', 'tif', 'exr', 'hdr', 'pic', 'dib', 'jpe', 'jp2', 'j2k', 'jpf', 'jpx', 'jpm', 'mj2', 'jxr', 'hdp', 'wdp', 'cur', 'ico', 'ani', 'icns', 'bpg', 'jp2', 'j2k', 'jpf', 'jpx', 'jpm', 'mj2', 'jxr', 'hdp', 'wdp', 'cur', 'ico', 'ani', 'icns', 'bpg')
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--source',    '-s',  type=str,                   help="Path to the source. Single image, video, directory of images, directory of videos is supported.")
+    parser.add_argument('--dest',      '-d',  type=str, default=None,     help="Path to destination. Results will be stored in current directory if not specified.")
+    parser.add_argument('--type',      '-t',  type=str, default='rgba',   help="Specify output type. If not specified, output results will make the background transparent. Please refer to the documentation for other types.")
+    parser.add_argument('--reverse',  '-R',  action='store_true',        help="Output will be reverse and foreground will be removed instead of background if specified.")
+    parser.add_argument('--format',    '-f',  type=str, default=None,     help="Specify output format. If not specified, it will be saved with the format of input.")
+    parser.add_argument('--resize',    '-r',  type=str, default='static', help="Specify resizing method. If not specified, static resize will be used. Choose from (static|dynamic).")
+    parser.add_argument('--jit',       '-j',  action='store_true',        help="Speed up inference speed by using torchscript, but decreases output quality.")
+    parser.add_argument('--device',    '-D',  type=str, default=None,     help="Designate device. If not specified, it will find available device.")
+    parser.add_argument('--mode',      '-m',  type=str, default='base',   help="choose between base and fast mode. Also, use base-nightly for nightly release checkpoint.")
+    parser.add_argument('--ckpt',      '-c',  type=str, default=None,     help="Designate checkpoint. If not specified, it will download or load pre-downloaded default checkpoint.")
+    parser.add_argument('--threshold', '-th', type=str, default=None,     help="Designate threshold. If specified, it will output hard prediction above threshold. If not specified, it will output soft prediction.")
+    return parser.parse_args()
+
+def get_backend():
+    if torch.cuda.is_available():
+        return "cuda:0"
+    elif torch.backends.mps.is_available():
+        return "mps:0"
+    else:
+        return "cpu"
+    
+def load_config(config_dir, easy=True):
+    cfg = yaml.load(open(config_dir), yaml.FullLoader)
+    if easy is True:
+        cfg = EasyDict(cfg)
+    return cfg
+
+def get_format(source):
+    img_count = len([i for i in source if i.lower().endswith(IMG_EXTS)])
+    vid_count = len([i for i in source if i.lower().endswith(VID_EXTS)])
+    
+    if img_count * vid_count != 0:
+        return ''
+    elif img_count != 0:
+        return 'Image'
+    elif vid_count != 0:
+        return 'Video'
+    else:
+        return ''
+
+def sort(x):
+    convert = lambda text: int(text) if text.isdigit() else text.lower()
+    alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]
+    return sorted(x, key=alphanum_key)
+
+def download_and_unzip(filename, url, dest, unzip=True, **kwargs):
+    if not os.path.isdir(dest):
+        os.makedirs(dest, exist_ok=True)
+    
+    if os.path.isfile(os.path.join(dest, filename)) is False:
+        os.system("wget -O {} {}".format(os.path.join(dest, filename), url))
+    elif 'md5' in kwargs.keys() and kwargs['md5'] != hashlib.md5(open(os.path.join(dest, filename), 'rb').read()).hexdigest():
+        os.system("wget -O {} {}".format(os.path.join(dest, filename), url))
+        
+    if unzip:
+        os.system("unzip -o {} -d {}".format(os.path.join(dest, filename), dest))
+        os.system("rm {}".format(os.path.join(dest, filename)))
+        
+class dynamic_resize:
+    def __init__(self, L=1280): 
+        self.L = L
+                    
+    def __call__(self, img):
+        size = list(img.size)
+        if (size[0] >= size[1]) and size[1] > self.L: 
+            size[0] = size[0] / (size[1] / self.L)
+            size[1] = self.L
+        elif (size[1] > size[0]) and size[0] > self.L:
+            size[1] = size[1] / (size[0] / self.L)
+            size[0] = self.L
+        size = (int(round(size[0] / 32)) * 32, int(round(size[1] / 32)) * 32)
+    
+        return img.resize(size, Image.BILINEAR)
+
+class dynamic_resize_a(ImageOnlyTransform):
+    def __init__(self, L=1280, always_apply=False, p=1.0):
+        super(dynamic_resize_a, self).__init__(always_apply, p)
+        self.L = L
+
+    def apply(self, img, **params):
+        size = list(img.shape[:2])
+        if (size[0] >= size[1]) and size[1] > self.L: 
+            size[0] = size[0] / (size[1] / self.L)
+            size[1] = self.L
+        elif (size[1] > size[0]) and size[0] > self.L:
+            size[1] = size[1] / (size[0] / self.L)
+            size[0] = self.L
+        size = (int(round(size[0] / 32)) * 32, int(round(size[1] / 32)) * 32)
+
+        return A.resize(img, height=size[0], width=size[1])
+
+    def get_transform_init_args_names(self):
+        return ("L",)
+    
+class static_resize:
+    def __init__(self, size=[1024, 1024]): 
+        self.size = size
+                    
+    def __call__(self, img):
+        return img.resize(self.size, Image.BILINEAR)    
+
+class normalize:
+    def __init__(self, mean=None, std=None, div=255):
+        self.mean = mean if mean is not None else 0.0
+        self.std = std if std is not None else 1.0
+        self.div = div
+        
+    def __call__(self, img):
+        img /= self.div
+        img -= self.mean
+        img /= self.std
+            
+        return img
+    
+class tonumpy:
+    def __init__(self):
+        pass
+
+    def __call__(self, img):
+        img = np.array(img, dtype=np.float32)
+        return img
+    
+class totensor:
+    def __init__(self):
+        pass
+
+    def __call__(self, img):
+        img = img.transpose((2, 0, 1))
+        img = torch.from_numpy(img).float()
+        
+        return img
+
+class ImageLoader:
+    def __init__(self, root):
+        if os.path.isdir(root):
+            self.images = [os.path.join(root, f) for f in os.listdir(root) if f.lower().endswith(('.jpg', '.png', '.jpeg'))]
+            self.images = sort(self.images)
+        elif os.path.isfile(root):
+            self.images = [root]
+        self.size = len(self.images)
+
+    def __iter__(self):
+        self.index = 0
+        return self
+
+    def __next__(self):
+        if self.index == self.size:
+            raise StopIteration
+        
+        img = Image.open(self.images[self.index]).convert('RGB')
+        name = os.path.split(self.images[self.index])[-1]
+        # name = os.path.splitext(name)[0]
+            
+        self.index += 1
+        return img, name
+
+    def __len__(self):
+        return self.size
+    
+class VideoLoader:
+    def __init__(self, root):
+        if os.path.isdir(root):
+            self.videos = [os.path.join(root, f) for f in os.listdir(root) if f.lower().endswith(('.mp4', '.avi', 'mov'))]
+        elif os.path.isfile(root):
+            self.videos = [root]
+        self.size = len(self.videos)
+
+    def __iter__(self):
+        self.index = 0
+        self.cap = None
+        self.fps = None
+        return self
+
+    def __next__(self):
+        if self.index == self.size:
+            raise StopIteration
+        
+        if self.cap is None:
+            self.cap = cv2.VideoCapture(self.videos[self.index])
+            self.fps = self.cap.get(cv2.CAP_PROP_FPS)
+        ret, frame = self.cap.read()
+        name = os.path.split(self.videos[self.index])[-1]
+        # name = os.path.splitext(name)[0]
+        if ret is False:
+            self.cap.release()
+            self.cap = None
+            img = None
+            self.index += 1
+        
+        else:
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            img = Image.fromarray(frame).convert('RGB')
+            
+        return img, name
+    
+    def __len__(self):
+        return self.size
+    
+class WebcamLoader:
+    def __init__(self, ID):
+        self.ID = int(ID)
+        self.cap = cv2.VideoCapture(self.ID)
+        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
+        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
+        self.imgs = []
+        self.imgs.append(self.cap.read()[1])
+        self.thread = Thread(target=self.update, daemon=True)
+        self.thread.start()
+        
+    def update(self):
+        while self.cap.isOpened():
+            ret, frame = self.cap.read()
+            if ret is True:
+                self.imgs.append(frame)
+            else:
+                break
+        
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if len(self.imgs) > 0:
+            frame = self.imgs[-1]
+        else:
+            frame = Image.fromarray(np.zeros((480, 640, 3)).astype(np.uint8))
+            
+        if self.thread.is_alive() is False or cv2.waitKey(1) == ord('q'):
+            cv2.destroyAllWindows()
+            raise StopIteration
+        
+        else:
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame = Image.fromarray(frame).convert('RGB')
+        
+        del self.imgs[:-1]
+        return frame, None
+
+    def __len__(self):
+        return 0