add app.py and related files

app.py

@@ -0,0 +1,137 @@
+import streamlit as st
+from diffusers import StableDiffusionInpaintPipeline
+import os
+
+from tqdm import tqdm
+from PIL import Image
+import numpy as np
+import cv2
+import warnings
+from huggingface_hub import hf_hub_download
+
+warnings.filterwarnings("ignore", category=FutureWarning)
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+
+from data.base_dataset import Normalize_image
+from utils.saving_utils import load_checkpoint_mgpu
+from networks import U2NET
+import argparse
+from enum import Enum
+from rembg import remove
+from dataclasses import dataclass
+
+
+@dataclass
+class StableFashionCLIArgs:
+    image
+    part
+    resolution
+    promt
+    num_steps
+    guidance_scale
+    rembg
+
+
+class Parts:
+    UPPER = 1
+    LOWER = 2
+
+
+def load_u2net():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    checkpoint_path = hf_hub_download(repo_id="maiti/cloth-segmentation", filename="cloth_segm_u2net_latest.pth")
+    net = U2NET(in_ch=3, out_ch=4)
+    net = load_checkpoint_mgpu(net, checkpoint_path)
+    net = net.to(device)
+    net = net.eval()
+    return net
+
+def change_bg_color(rgba_image, color):
+    new_image = Image.new("RGBA", rgba_image.size, color)
+    new_image.paste(rgba_image, (0, 0), rgba_image)
+    return new_image.convert("RGB")
+
+def load_inpainting_pipeline():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    inpainting_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
+            "runwayml/stable-diffusion-inpainting",
+            revision="fp16",
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        ).to(device)
+    return inpainting_pipeline
+
+
+def process_image(args, inpainting_pipeline, net):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    image_path = args.image
+    transforms_list = []
+    transforms_list += [transforms.ToTensor()]
+    transforms_list += [Normalize_image(0.5, 0.5)]
+    transform_rgb = transforms.Compose(transforms_list)
+    img = Image.open(image_path)
+    img = img.convert("RGB")
+    img = img.resize((args.resolution, args.resolution))
+    if args.rembg:
+        img_with_green_bg = remove(img)
+        img_with_green_bg = change_bg_color(img_with_green_bg, color="GREEN")
+        img_with_green_bg = img_with_green_bg.convert("RGB")
+    else:
+        img_with_green_bg = img
+    image_tensor = transform_rgb(img_with_green_bg)
+    image_tensor = image_tensor.unsqueeze(0)
+    output_tensor = net(image_tensor.to(device))
+    output_tensor = F.log_softmax(output_tensor[0], dim=1)
+    output_tensor = torch.max(output_tensor, dim=1, keepdim=True)[1]
+    output_tensor = torch.squeeze(output_tensor, dim=0)
+    output_tensor = torch.squeeze(output_tensor, dim=0)
+    output_arr = output_tensor.cpu().numpy()
+    mask_code = eval(f"Parts.{args.part.upper()}")
+    mask = (output_arr == mask_code)
+    output_arr[mask] = 1
+    output_arr[~mask] = 0
+    output_arr *= 255
+    mask_PIL = Image.fromarray(output_arr.astype("uint8"), mode="L")
+    clothed_image_from_pipeline = inpainting_pipeline(prompt=args.prompt,
+                                                    image=img_with_green_bg,
+                                                    mask_image=mask_PIL,
+                                                    width=args.resolution,
+                                                    height=args.resolution,
+                                                    guidance_scale=args.guidance_scale,
+                                                    num_inference_steps=args.num_steps).images[0]
+    clothed_image_from_pipeline = remove(clothed_image_from_pipeline)
+    clothed_image_from_pipeline = change_bg_color(clothed_image_from_pipeline, "WHITE")
+    return clothed_image_from_pipeline.convert("RGB")
+
+
+st.title("Stable Fashion Huggingface Spaces")
+file_name = st.file_uploader("Upload a clear full length picture of yourself, preferably in a less noisy background")
+net = load_u2net()
+inpainting_pipeline = load_inpainting_pipeline()
+
+if file_name is not None:
+    image = Image.open(file_name)
+    stable_fashion_args = StableFashionCLIArgs()
+    stable_fashion_args.image = image
+    body_part = st.radio("Would you like to try clothes on your upper body (such as shirts, kurtas etc) or lower (Jeans, Pants etc)? ", ('Upper', 'Lower'))
+    stable_fashion_args.part = body_part
+    resolution = st.radio("Which resolution would you like to get the resulting picture in? (Keep in mind, higher the resolution, higher the queue times)", (128, 256, 512))
+    stable_fashion_args.resolution = resolution
+    rembg_status = st.radio("Would you like to remove background in your image before putting new clothes on you? (Sometimes it results in better images)", ("Yes", "No"))
+    stable_fashion_args.rembg = (rembg_status == "Yes")
+    guidance_scale = st.slider("Select a guidance scale. 7.5 gives the best results.", 1.0, 15.0, value=7.5)
+    stable_fashion_args.guidance_scale = guidance_scale
+    prompt = st.text_input('Write the description of cloth you want to try', 'a bright yellow t shirt')
+    stable_fashion_args.prompt = guidance_scale
+
+    num_steps = st.slider("No. of inference steps for the diffusion process", 5, 50, value=25)
+
+
+    result_image = process_image(stable_fashion_args, inpainting_pipeline, net)
+    st.image(result_image, caption='Sunrise by the mountains')
+
+
+

data/base_dataset.py

@@ -0,0 +1,189 @@
+import os
+from PIL import Image
+import cv2
+import numpy as np
+import random
+
+import torch
+import torch.utils.data as data
+import torchvision.transforms as transforms
+
+
+class BaseDataset(data.Dataset):
+    def __init__(self):
+        super(BaseDataset, self).__init__()
+
+    def name(self):
+        return "BaseDataset"
+
+    def initialize(self, opt):
+        pass
+
+
+class Rescale_fixed(object):
+    """Rescale the input image into given size.
+
+    Args:
+        (w,h) (tuple): output size or x (int) then resized will be done in (x,x).
+    """
+
+    def __init__(self, output_size):
+        self.output_size = output_size
+
+    def __call__(self, image):
+        return image.resize(self.output_size, Image.BICUBIC)
+
+
+class Rescale_custom(object):
+    """Rescale the input image and target image into randomly selected size with lower bound of min_size arg.
+
+    Args:
+        min_size (int): Minimum desired output size.
+    """
+
+    def __init__(self, min_size, max_size):
+        assert isinstance(min_size, (int, float))
+        self.min_size = min_size
+        self.max_size = max_size
+
+    def __call__(self, sample):
+
+        input_image, target_image = sample["input_image"], sample["target_image"]
+
+        assert input_image.size == target_image.size
+        w, h = input_image.size
+
+        # Randomly select size to resize
+        if min(self.max_size, h, w) > self.min_size:
+            self.output_size = np.random.randint(
+                self.min_size, min(self.max_size, h, w)
+            )
+        else:
+            self.output_size = self.min_size
+
+        # calculate new size by keeping aspect ratio same
+        if h > w:
+            new_h, new_w = self.output_size * h / w, self.output_size
+        else:
+            new_h, new_w = self.output_size, self.output_size * w / h
+
+        new_w, new_h = int(new_w), int(new_h)
+        input_image = input_image.resize((new_w, new_h), Image.BICUBIC)
+        target_image = target_image.resize((new_w, new_h), Image.BICUBIC)
+        return {"input_image": input_image, "target_image": target_image}
+
+
+class ToTensor(object):
+    """Convert ndarrays in sample to Tensors."""
+
+    def __init__(self):
+        self.totensor = transforms.ToTensor()
+
+    def __call__(self, sample):
+        input_image, target_image = sample["input_image"], sample["target_image"]
+
+        return {
+            "input_image": self.totensor(input_image),
+            "target_image": self.totensor(target_image),
+        }
+
+
+class RandomCrop_custom(object):
+    """Crop randomly the image in a sample.
+
+    Args:
+        output_size (tuple or int): Desired output size. If int, square crop
+            is made.
+    """
+
+    def __init__(self, output_size):
+        assert isinstance(output_size, (int, tuple))
+        if isinstance(output_size, int):
+            self.output_size = (output_size, output_size)
+        else:
+            assert len(output_size) == 2
+            self.output_size = output_size
+
+        self.randomcrop = transforms.RandomCrop(self.output_size)
+
+    def __call__(self, sample):
+        input_image, target_image = sample["input_image"], sample["target_image"]
+        cropped_imgs = self.randomcrop(torch.cat((input_image, target_image)))
+
+        return {
+            "input_image": cropped_imgs[
+                :3,
+                :,
+            ],
+            "target_image": cropped_imgs[
+                3:,
+                :,
+            ],
+        }
+
+
+class Normalize_custom(object):
+    """Normalize given dict into given mean and standard dev
+
+    Args:
+        mean (tuple or int): Desired mean to substract from dict's tensors
+        std (tuple or int): Desired std to divide from dict's tensors
+    """
+
+    def __init__(self, mean, std):
+        assert isinstance(mean, (float, tuple))
+        if isinstance(mean, float):
+            self.mean = (mean, mean, mean)
+        else:
+            assert len(mean) == 3
+            self.mean = mean
+
+        if isinstance(std, float):
+            self.std = (std, std, std)
+        else:
+            assert len(std) == 3
+            self.std = std
+
+        self.normalize = transforms.Normalize(self.mean, self.std)
+
+    def __call__(self, sample):
+        input_image, target_image = sample["input_image"], sample["target_image"]
+
+        return {
+            "input_image": self.normalize(input_image),
+            "target_image": self.normalize(target_image),
+        }
+
+
+class Normalize_image(object):
+    """Normalize given tensor into given mean and standard dev
+
+    Args:
+        mean (float): Desired mean to substract from tensors
+        std (float): Desired std to divide from tensors
+    """
+
+    def __init__(self, mean, std):
+        assert isinstance(mean, (float))
+        if isinstance(mean, float):
+            self.mean = mean
+
+        if isinstance(std, float):
+            self.std = std
+
+        self.normalize_1 = transforms.Normalize(self.mean, self.std)
+        self.normalize_3 = transforms.Normalize([self.mean] * 3, [self.std] * 3)
+        self.normalize_18 = transforms.Normalize([self.mean] * 18, [self.std] * 18)
+
+    def __call__(self, image_tensor):
+        if image_tensor.shape[0] == 1:
+            return self.normalize_1(image_tensor)
+
+        elif image_tensor.shape[0] == 3:
+            return self.normalize_3(image_tensor)
+
+        elif image_tensor.shape[0] == 18:
+            return self.normalize_18(image_tensor)
+
+        else:
+            assert "Please set proper channels! Normlization implemented only for 1, 3 and 18"

main.py

@@ -0,0 +1,112 @@
+from diffusers import StableDiffusionInpaintPipeline
+import os
+
+from tqdm import tqdm
+from PIL import Image
+import numpy as np
+import cv2
+import warnings
+
+warnings.filterwarnings("ignore", category=FutureWarning)
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+
+from data.base_dataset import Normalize_image
+from utils.saving_utils import load_checkpoint_mgpu
+from networks import U2NET
+import argparse
+from enum import Enum
+from rembg import remove
+
+class Parts:
+    UPPER = 1
+    LOWER = 2
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(
+        description="Stable Fashion API, allows you to picture yourself in any cloth your imagination can think of!"
+    )
+    parser.add_argument('--image', type=str, required=True, help='path to image')
+    parser.add_argument('--part', choices=['upper', 'lower'], default='upper', type=str)
+    parser.add_argument('--resolution', choices=[256, 512, 1024, 2048], default=256, type=int)
+    parser.add_argument('--prompt', type=str, default="A pink cloth")
+    parser.add_argument('--num_steps', type=int, default=5)
+    parser.add_argument('--guidance_scale', type=float, default=7.5)
+    parser.add_argument('--rembg', action='store_true')
+    parser.add_argument('--output', default='output.jpg', type=str)
+    args, _ = parser.parse_known_args()
+    return args
+
+
+def load_u2net():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    checkpoint_path = os.path.join("trained_checkpoint", "cloth_segm_u2net_latest.pth")
+    net = U2NET(in_ch=3, out_ch=4)
+    net = load_checkpoint_mgpu(net, checkpoint_path)
+    net = net.to(device)
+    net = net.eval()
+    return net
+
+def change_bg_color(rgba_image, color):
+    new_image = Image.new("RGBA", rgba_image.size, color)
+    new_image.paste(rgba_image, (0, 0), rgba_image)
+    return new_image.convert("RGB")
+
+
+def load_inpainting_pipeline():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    inpainting_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
+            "runwayml/stable-diffusion-inpainting",
+            revision="fp16",
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        ).to(device)
+    return inpainting_pipeline
+def process_image(args, inpainting_pipeline, net):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    image_path = args.image
+    transforms_list = []
+    transforms_list += [transforms.ToTensor()]
+    transforms_list += [Normalize_image(0.5, 0.5)]
+    transform_rgb = transforms.Compose(transforms_list)
+    img = Image.open(image_path)
+    img = img.convert("RGB")
+    img = img.resize((args.resolution, args.resolution))
+    if args.rembg:
+        img_with_green_bg = remove(img)
+        img_with_green_bg = change_bg_color(img_with_green_bg, color="GREEN")
+        img_with_green_bg = img_with_green_bg.convert("RGB")
+    else:
+        img_with_green_bg = img
+    image_tensor = transform_rgb(img_with_green_bg)
+    image_tensor = image_tensor.unsqueeze(0)
+    output_tensor = net(image_tensor.to(device))
+    output_tensor = F.log_softmax(output_tensor[0], dim=1)
+    output_tensor = torch.max(output_tensor, dim=1, keepdim=True)[1]
+    output_tensor = torch.squeeze(output_tensor, dim=0)
+    output_tensor = torch.squeeze(output_tensor, dim=0)
+    output_arr = output_tensor.cpu().numpy()
+    mask_code = eval(f"Parts.{args.part.upper()}")
+    mask = (output_arr == mask_code)
+    output_arr[mask] = 1
+    output_arr[~mask] = 0
+    output_arr *= 255
+    mask_PIL = Image.fromarray(output_arr.astype("uint8"), mode="L")
+    clothed_image_from_pipeline = inpainting_pipeline(prompt=args.prompt,
+                                                    image=img_with_green_bg,
+                                                    mask_image=mask_PIL,
+                                                    width=args.resolution,
+                                                    height=args.resolution,
+                                                    guidance_scale=args.guidance_scale,
+                                                    num_inference_steps=args.num_steps).images[0]
+    clothed_image_from_pipeline = remove(clothed_image_from_pipeline)
+    clothed_image_from_pipeline = change_bg_color(clothed_image_from_pipeline, "WHITE")
+    return clothed_image_from_pipeline.convert("RGB")
+if __name__ == '__main__':
+    args = parse_arguments()
+    net = load_u2net()
+    inpainting_pipeline = load_inpainting_pipeline()
+    result_image = process_image(args, inpainting_pipeline, net)
+    result_image.save(args.output)

networks/__init__.py

@@ -0,0 +1 @@
+from .u2net import U2NET

networks/u2net.py

@@ -0,0 +1,565 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class REBNCONV(nn.Module):
+    def __init__(self, in_ch=3, out_ch=3, dirate=1):
+        super(REBNCONV, self).__init__()
+
+        self.conv_s1 = nn.Conv2d(
+            in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate
+        )
+        self.bn_s1 = nn.BatchNorm2d(out_ch)
+        self.relu_s1 = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+
+        hx = x
+        xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
+
+        return xout
+
+
+## upsample tensor 'src' to have the same spatial size with tensor 'tar'
+def _upsample_like(src, tar):
+
+    src = F.upsample(src, size=tar.shape[2:], mode="bilinear")
+
+    return src
+
+
+### RSU-7 ###
+class RSU7(nn.Module):  # UNet07DRES(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU7, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool5 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv7 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv6d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+
+        hx = x
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+        hx = self.pool5(hx5)
+
+        hx6 = self.rebnconv6(hx)
+
+        hx7 = self.rebnconv7(hx6)
+
+        hx6d = self.rebnconv6d(torch.cat((hx7, hx6), 1))
+        hx6dup = _upsample_like(hx6d, hx5)
+
+        hx5d = self.rebnconv5d(torch.cat((hx6dup, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        """
+        del hx1, hx2, hx3, hx4, hx5, hx6, hx7
+        del hx6d, hx5d, hx3d, hx2d
+        del hx2dup, hx3dup, hx4dup, hx5dup, hx6dup
+        """
+
+        return hx1d + hxin
+
+
+### RSU-6 ###
+class RSU6(nn.Module):  # UNet06DRES(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU6, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+
+        hx6 = self.rebnconv6(hx5)
+
+        hx5d = self.rebnconv5d(torch.cat((hx6, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        """
+        del hx1, hx2, hx3, hx4, hx5, hx6
+        del hx5d, hx4d, hx3d, hx2d
+        del hx2dup, hx3dup, hx4dup, hx5dup
+        """
+
+        return hx1d + hxin
+
+
+### RSU-5 ###
+class RSU5(nn.Module):  # UNet05DRES(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU5, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+
+        hx5 = self.rebnconv5(hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        """
+        del hx1, hx2, hx3, hx4, hx5
+        del hx4d, hx3d, hx2d
+        del hx2dup, hx3dup, hx4dup
+        """
+
+        return hx1d + hxin
+
+
+### RSU-4 ###
+class RSU4(nn.Module):  # UNet04DRES(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        """
+        del hx1, hx2, hx3, hx4
+        del hx3d, hx2d
+        del hx2dup, hx3dup
+        """
+
+        return hx1d + hxin
+
+
+### RSU-4F ###
+class RSU4F(nn.Module):  # UNet04FRES(nn.Module):
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4F, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=2)
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=4)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=8)
+
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=4)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=2)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx2 = self.rebnconv2(hx1)
+        hx3 = self.rebnconv3(hx2)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))
+        hx2d = self.rebnconv2d(torch.cat((hx3d, hx2), 1))
+        hx1d = self.rebnconv1d(torch.cat((hx2d, hx1), 1))
+
+        """
+        del hx1, hx2, hx3, hx4
+        del hx3d, hx2d
+        """
+
+        return hx1d + hxin
+
+
+##### U^2-Net ####
+class U2NET(nn.Module):
+    def __init__(self, in_ch=3, out_ch=1):
+        super(U2NET, self).__init__()
+
+        self.stage1 = RSU7(in_ch, 32, 64)
+        self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage2 = RSU6(64, 32, 128)
+        self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage3 = RSU5(128, 64, 256)
+        self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage4 = RSU4(256, 128, 512)
+        self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage5 = RSU4F(512, 256, 512)
+        self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage6 = RSU4F(512, 256, 512)
+
+        # decoder
+        self.stage5d = RSU4F(1024, 256, 512)
+        self.stage4d = RSU4(1024, 128, 256)
+        self.stage3d = RSU5(512, 64, 128)
+        self.stage2d = RSU6(256, 32, 64)
+        self.stage1d = RSU7(128, 16, 64)
+
+        self.side1 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side2 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side3 = nn.Conv2d(128, out_ch, 3, padding=1)
+        self.side4 = nn.Conv2d(256, out_ch, 3, padding=1)
+        self.side5 = nn.Conv2d(512, out_ch, 3, padding=1)
+        self.side6 = nn.Conv2d(512, out_ch, 3, padding=1)
+
+        self.outconv = nn.Conv2d(6 * out_ch, out_ch, 1)
+
+    def forward(self, x):
+
+        hx = x
+
+        # stage 1
+        hx1 = self.stage1(hx)
+        hx = self.pool12(hx1)
+
+        # stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        # stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        # stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        # stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        # stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6, hx5)
+
+        # -------------------- decoder --------------------
+        hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))
+
+        # side output
+        d1 = self.side1(hx1d)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2, d1)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3, d1)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4, d1)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5, d1)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6, d1)
+
+        d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1))
+
+        """
+        del hx1, hx2, hx3, hx4, hx5, hx6
+        del hx5d, hx4d, hx3d, hx2d, hx1d
+        del hx6up, hx5dup, hx4dup, hx3dup, hx2dup
+        """
+
+        return d0, d1, d2, d3, d4, d5, d6
+
+
+### U^2-Net small ###
+class U2NETP(nn.Module):
+    def __init__(self, in_ch=3, out_ch=1):
+        super(U2NETP, self).__init__()
+
+        self.stage1 = RSU7(in_ch, 16, 64)
+        self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage2 = RSU6(64, 16, 64)
+        self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage3 = RSU5(64, 16, 64)
+        self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage4 = RSU4(64, 16, 64)
+        self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage5 = RSU4F(64, 16, 64)
+        self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage6 = RSU4F(64, 16, 64)
+
+        # decoder
+        self.stage5d = RSU4F(128, 16, 64)
+        self.stage4d = RSU4(128, 16, 64)
+        self.stage3d = RSU5(128, 16, 64)
+        self.stage2d = RSU6(128, 16, 64)
+        self.stage1d = RSU7(128, 16, 64)
+
+        self.side1 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side2 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side3 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side4 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side5 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side6 = nn.Conv2d(64, out_ch, 3, padding=1)
+
+        self.outconv = nn.Conv2d(6 * out_ch, out_ch, 1)
+
+    def forward(self, x):
+
+        hx = x
+
+        # stage 1
+        hx1 = self.stage1(hx)
+        hx = self.pool12(hx1)
+
+        # stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        # stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        # stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        # stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        # stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6, hx5)
+
+        # decoder
+        hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))
+
+        # side output
+        d1 = self.side1(hx1d)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2, d1)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3, d1)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4, d1)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5, d1)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6, d1)
+
+        d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1))
+
+        """
+        del hx1, hx2, hx3, hx4, hx5, hx6
+        del hx5d, hx4d, hx3d, hx2d, hx1d
+        del hx6up, hx5dup, hx4dup, hx3dup, hx2dup
+        """
+
+        return d0, d1, d2, d3, d4, d5, d6

requirements.txt

@@ -0,0 +1,91 @@
+absl-py==1.2.0
+accelerate==0.12.0
+aiohttp==3.8.1
+aiosignal==1.2.0
+asttokens==2.0.8
+async-timeout==4.0.2
+attrs==22.1.0
+backcall==0.2.0
+beautifulsoup4==4.11.1
+bitsandbytes==0.33.1
+cachetools==5.2.0
+charset-normalizer==2.1.1
+clip @ git+https://github.com/openai/CLIP.git@d50d76daa670286dd6cacf3bcd80b5e4823fc8e1
+datasets==2.4.0
+decorator==5.1.1
+diffusers @ git+https://github.com/ovshake/diffusers@5bbecb751764248755943b57d900ae14a7f43a75
+dill==0.3.5.1
+executing==1.0.0
+filelock==3.8.0
+frozenlist==1.3.1
+fsspec==2022.8.2
+ftfy==6.1.1
+gdown==4.5.1
+google-auth==2.11.0
+google-auth-oauthlib==0.4.6
+grpcio==1.47.0
+huggingface-hub==0.11.0
+idna==3.3
+importlib-metadata==4.12.0
+ipdb==0.13.9
+ipython==8.4.0
+jedi==0.18.1
+Jinja2==3.1.2
+joblib==1.1.0
+Markdown==3.4.1
+MarkupSafe==2.1.1
+matplotlib-inline==0.1.6
+modelcards==0.1.6
+multidict==6.0.2
+multiprocess==0.70.13
+numpy==1.23.2
+oauthlib==3.2.0
+opencv-python==4.6.0.66
+packaging==21.3
+pandas==1.4.4
+parso==0.8.3
+pexpect==4.8.0
+pickleshare==0.7.5
+Pillow==9.2.0
+prompt-toolkit==3.0.30
+protobuf==3.19.4
+psutil==5.9.1
+ptyprocess==0.7.0
+pure-eval==0.2.2
+pyarrow==9.0.0
+pyasn1==0.4.8
+pyasn1-modules==0.2.8
+Pygments==2.13.0
+pyparsing==3.0.9
+PySocks==1.7.1
+python-dateutil==2.8.2
+pytz==2022.2.1
+PyYAML==6.0
+regex==2022.8.17
+requests==2.28.1
+requests-oauthlib==1.3.1
+responses==0.18.0
+rsa==4.9
+six==1.16.0
+soupsieve==2.3.2.post1
+stack-data==0.5.0
+tensorboard==2.10.0
+tensorboard-data-server==0.6.1
+tensorboard-plugin-wit==1.8.1
+tensorboardX==2.5.1
+tokenizers==0.13.2
+toml==0.10.2
+torch==1.13.0+cu116
+torchaudio==0.13.0+cu116
+torchvision==0.14.0+cu116
+tqdm==4.64.0
+traitlets==5.3.0
+transformers==4.24.0
+typing_extensions==4.3.0
+urllib3==1.26.12
+wcwidth==0.2.5
+Werkzeug==2.2.2
+xxhash==3.0.0
+yarl==1.8.1
+zipp==3.8.1
+rembg

utils/saving_utils.py

@@ -0,0 +1,45 @@
+import os
+import copy
+import cv2
+import numpy as np
+from collections import OrderedDict
+
+import torch
+
+
+def load_checkpoint(model, checkpoint_path):
+    if not os.path.exists(checkpoint_path):
+        print("----No checkpoints at given path----")
+        return
+    model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device("cpu")))
+    print("----checkpoints loaded from path: {}----".format(checkpoint_path))
+    return model
+
+
+def load_checkpoint_mgpu(model, checkpoint_path):
+    if not os.path.exists(checkpoint_path):
+        print("----No checkpoints at given path----")
+        return
+    model_state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))
+    new_state_dict = OrderedDict()
+    for k, v in model_state_dict.items():
+        name = k[7:]  # remove `module.`
+        new_state_dict[name] = v
+
+    model.load_state_dict(new_state_dict)
+    print("----checkpoints loaded from path: {}----".format(checkpoint_path))
+    return model
+
+
+def save_checkpoint(model, save_path):
+    print(save_path)
+    if not os.path.exists(os.path.dirname(save_path)):
+        os.makedirs(os.path.dirname(save_path))
+    torch.save(model.state_dict(), save_path)
+
+
+def save_checkpoints(opt, itr, net):
+    save_checkpoint(
+        net,
+        os.path.join(opt.save_dir, "checkpoints", "itr_{:08d}_u2net.pth".format(itr)),
+    )