preprocess garment added

.gitignore

@@ -0,0 +1,8 @@
+.env
+.DS_Store
+input_image
+output_image
+cloth-mask
+__pycache__
+*.pyc
+venv

app.py

@@ -0,0 +1,82 @@
+from dotenv import load_dotenv
+load_dotenv()
+
+import glob
+import os
+from PIL import Image
+
+import gradio as gr
+import preprocess
+from huggingface_hub import login
+
+def extract_garment(input_img, cls):
+    print(input_img, type(input_img), cls)
+
+    input_dir = "input_image"
+    output_dir = "output_image"
+
+    os.makedirs(input_dir, exist_ok=True)
+    os.makedirs(output_dir, exist_ok=True)
+
+    for f in glob.glob(input_dir + "/*.*"):
+        os.remove(f)
+
+    for f in glob.glob(output_dir + "/*.*"):
+        os.remove(f)
+
+    for f in glob.glob("cloth-mask/*.*"):
+        os.remove(f)
+
+    input_img.save(os.path.join(input_dir, "img.jpg"))
+
+    preprocess.extract_garment(inputs_dir=input_dir, outputs_dir=output_dir, cls=cls)
+
+    return Image.open(glob.glob(output_dir + "/*.*")[0])
+
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 720px;
+}
+"""
+
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(f"""
+        # Clothes Extraction using U2Net
+        Pull out clothes like tops, bottoms, and dresses from a photo. This implementation is based on the [U2Net](https://github.com/xuebinqin/U-2-Net) model.
+        """)
+
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(label="Input Image", type='pil', height="400px", show_label=True)
+                dropdown = gr.Dropdown(["upper", "lower", "dress"], value="upper", label="Extract garment",
+                                       info="Select the garment type you wish to extract!")
+
+            output_image = gr.Image(label="Extracted garment", type='pil', height="400px", show_label=True,
+                                    show_download_button=True)
+
+        with gr.Row():
+            submit_button = gr.Button("Submit", variant='primary', scale=1)
+            reset_button = gr.ClearButton(value="Reset", scale=1)
+
+    gr.on(
+        triggers=[submit_button.click],
+        fn=extract_garment,
+        inputs=[input_image, dropdown],
+        outputs=[output_image]
+    )
+
+    reset_button.click(
+        fn=lambda: (None, "upper", None),
+        inputs=[],
+        outputs=[input_image, dropdown, output_image],
+        concurrency_limit=1,
+    )
+
+if __name__ == '__main__':
+    # login to hugging face
+    login(os.environ.get("HF_TOKEN"))
+    
+    demo.launch(show_api=True)

preprocess/__init__.py

@@ -0,0 +1 @@
+from .preprocess_garment import segment_garment, extract_garment

preprocess/load_u2net.py

@@ -0,0 +1,27 @@
+import os
+from collections import OrderedDict
+
+import torch
+
+from .u2net_cloth_segm import U2NET
+
+
+def load_cloth_segm_model(device, checkpoint_path, in_ch=3, out_ch=1):
+    if not os.path.exists(checkpoint_path):
+        print("Invalid path")
+        return
+
+    model = U2NET(in_ch=in_ch, out_ch=out_ch)
+
+    model_state_dict = torch.load(checkpoint_path, map_location=device)
+    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)
+    model = model.to(device=device)
+
+    print("Checkpoints loaded from path: {}".format(checkpoint_path))
+
+    return model

preprocess/preprocess_garment.py

@@ -0,0 +1,114 @@
+import glob
+import os
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from torchvision import transforms
+from tqdm import tqdm
+import joblib
+from huggingface_hub import hf_hub_download
+
+from .load_u2net import load_cloth_segm_model
+from .utils import NormalizeImage, naive_cutout, resize_by_bigger_index, image_resize
+
+
+def segment_garment(inputs_dir, outputs_dir, cls="all"):
+    os.makedirs(outputs_dir, exist_ok=True)
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    transform_fn = transforms.Compose(
+        [transforms.ToTensor(),
+         NormalizeImage(0.5, 0.5)]
+    )
+
+    # load model from huggingface
+    file_path = hf_hub_download(repo_id="tryonlabs/u2net-cloth-segmentation", filename="u2net_cloth_segm.pth")
+    
+    print("model loaded from huggingface:", file_path)
+    
+    net = load_cloth_segm_model(device, file_path, in_ch=3, out_ch=4)
+
+    images_list = sorted(os.listdir(inputs_dir))
+    pbar = tqdm(total=len(images_list))
+
+    for image_name in images_list:
+        img = Image.open(os.path.join(inputs_dir, image_name)).convert('RGB')
+        img_size = img.size
+        img = img.resize((768, 768), Image.BICUBIC)
+        image_tensor = transform_fn(img)
+        image_tensor = torch.unsqueeze(image_tensor, 0)
+
+        with torch.no_grad():
+            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_arr = output_tensor.cpu().numpy()
+
+        if cls == "all":
+            classes_to_save = []
+
+            # Check which classes are present in the image
+            for cls in range(1, 4):  # Exclude background class (0)
+                if np.any(output_arr == cls):
+                    classes_to_save.append(cls)
+        elif cls == "upper":
+            classes_to_save = [1]
+        elif cls == "lower":
+            classes_to_save = [2]
+        elif cls == "dress":
+            classes_to_save = [3]
+        else:
+            raise ValueError(f"Unknown cls: {cls}")
+
+        for cls1 in classes_to_save:
+            alpha_mask = (output_arr == cls1).astype(np.uint8) * 255
+            alpha_mask = alpha_mask[0]  # Selecting the first channel to make it 2D
+            alpha_mask_img = Image.fromarray(alpha_mask, mode='L')
+            alpha_mask_img = alpha_mask_img.resize(img_size, Image.BICUBIC)
+            alpha_mask_img.save(os.path.join(outputs_dir, f'{image_name.split(".")[0]}_{cls1}.jpg'))
+
+        pbar.update(1)
+
+    pbar.close()
+
+
+def extract_garment(inputs_dir, outputs_dir, cls="all", resize_to_width=None):
+    os.makedirs(outputs_dir, exist_ok=True)
+    cloth_mask_dir = os.path.join(Path(outputs_dir).parent.absolute(), "cloth-mask")
+    os.makedirs(cloth_mask_dir, exist_ok=True)
+
+    segment_garment(inputs_dir, os.path.join(Path(outputs_dir).parent.absolute(), "cloth-mask"), cls=cls)
+
+    images_path = sorted(glob.glob(os.path.join(inputs_dir, "*")))
+    masks_path = sorted(glob.glob(os.path.join(cloth_mask_dir, "*")))
+    img = Image.open(images_path[0])
+
+    for mask_path in masks_path:
+        mask = Image.open(mask_path)
+
+        cutout = np.array(naive_cutout(img, mask))
+        cutout = resize_by_bigger_index(cutout)
+
+        canvas = np.ones((1024, 768, 3), np.uint8) * 255
+        y1, y2 = (canvas.shape[0] - cutout.shape[0]) // 2, (canvas.shape[0] + cutout.shape[0]) // 2
+        x1, x2 = (canvas.shape[1] - cutout.shape[1]) // 2, (canvas.shape[1] + cutout.shape[1]) // 2
+
+        alpha_s = cutout[:, :, 3] / 255.0
+        alpha_l = 1.0 - alpha_s
+
+        for c in range(0, 3):
+            canvas[y1:y2, x1:x2, c] = (alpha_s * cutout[:, :, c] +
+                                       alpha_l * canvas[y1:y2, x1:x2, c])
+
+        # resize image before saving
+        if resize_to_width:
+            canvas = image_resize(canvas, width=resize_to_width)
+
+        canvas = Image.fromarray(canvas)
+
+        canvas.save(os.path.join(outputs_dir, f"{os.path.basename(mask_path).split('.')[0]}.jpg"), format='JPEG')

preprocess/u2net_cloth_segm.py

@@ -0,0 +1,550 @@
+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))
+
+        return d0, d1, d2, d3, d4, d5, d6

preprocess/utils.py

@@ -0,0 +1,91 @@
+import os
+from pathlib import Path
+
+import cv2
+from PIL import Image
+from torchvision import transforms
+
+
+class NormalizeImage(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! Normalization implemented only for 1, 3 and 18"
+
+
+def naive_cutout(img, mask):
+    empty = Image.new("RGBA", (img.size), 0)
+    cutout = Image.composite(img, empty, mask.resize(img.size, Image.LANCZOS))
+    return cutout
+
+
+def resize_by_bigger_index(crop):
+    # function resizes and keeps the aspect ratio same
+    crop_shape = crop.shape  # hxwxc
+    if crop_shape[0] / crop_shape[1] <= 1.33:
+        resized_crop = image_resize(crop, width=768)
+    else:
+        resized_crop = image_resize(crop, height=1024)
+    return resized_crop
+
+
+def image_resize(image, width=None, height=None):
+    dim = None
+    (h, w) = image.shape[:2]
+
+    if width is None and height is None:
+        return image
+
+    if width is None:
+        r = height / float(h)
+        dim = (int(w * r), height)
+
+    else:
+        r = width / float(w)
+        dim = (width, int(h * r))
+
+    resized = cv2.resize(image, dim)
+
+    return resized
+
+
+def convert_to_jpg(image_path, output_dir, size=None):
+    """
+    Convert image to jpg format
+    :param image_path: image path
+    :param output_dir: output directory
+    :param size: desired size of the image (w, h)
+    """
+    img = cv2.imread(image_path)
+    if size is not None:
+        img = image_resize(img, width=size[0], height=size[1])
+
+    filename = Path(image_path).name
+    cv2.imwrite(os.path.join(output_dir, filename.split(".")[0] + ".jpg"), img)

requirements.txt

@@ -0,0 +1,10 @@
+pillow
+gradio==4.44.1
+torch
+torchvision
+numpy==1.26.1
+tqdm
+opencv-python
+joblib
+huggingface-hub==0.25.0
+python-dotenv