Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,162 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+.idea/
+.vscode
+.DS_Store

README.md

@@ -1,12 +1,50 @@
----
-title: Vrtclothes
-emoji: 💻
-colorFrom: red
-colorTo: blue
-sdk: gradio
-sdk_version: 4.44.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Cloths Virtual Try On
+[![Open in colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/SwayamInSync/clothes-virtual-try-on/blob/main/setup_gradio.ipynb)
+
+## Updates
+- **[19/02/2024] From now on this repo won't receive any future updates from my side (Spoiler: It's not gone for good 😉. Expect its return, stronger than ever.) (Community Contributions & Issues discussions are still welcome 🤗)**
+- [26/12/2023] Added the Gradio interface and removed all the external depenency
+- [19/12/2023] Fixed the `openpose` installation and missing model weights issue
+- [19/12/2023] Replaced the `remove.bg` dependecy with `rembg`
+- [26/04/2023] Fixed the GAN generation issue
+
+## Star History
+[![Star History Chart](https://api.star-history.com/svg?repos=SwayamInSync/clothes-virtual-try-on&type=Date)](https://star-history.com/#SwayamInSync/clothes-virtual-try-on&Date)
+
+## Table of contents
+- [Cloths Virtual Try On](#cloths-virtual-try-on)
+  - [Table of contents](#table-of-contents)
+  - [General info](#general-info)
+  - [Demo](#demo)
+  - [Block Diagram](#block-diagram)
+  - [Methodology](#methodology)
+  - [Usage](#usage)
+  - [Citation](#citation)
+
+## General info
+
+This project is a part of a crework community project. While buying clothes online, it is difficult for a customer to select a desirable outfit in the first attempt because they can’t try on clothes. This project aims to solve this problem.
+
+<img width="383" alt="general_info" src="https://user-images.githubusercontent.com/63489382/163923011-c2898812-2491-4ec2-beb7-dcaaaf680e4f.png">
+
+
+## Demo
+
+https://user-images.githubusercontent.com/63489382/163922795-5dbb0f52-95e4-42c6-95d7-2d965abeba6d.mp4
+
+
+
+## Block Diagram
+![block_diagram_whole](https://user-images.githubusercontent.com/63489382/163922947-c1677f79-ad6f-4550-affc-7d4e80f0d247.png)
+
+
+## Methodology
+![block_diagram_detailed](https://user-images.githubusercontent.com/63489382/163922991-86d148c2-1a97-48a5-b4ec-d8c16819374a.png)
+
+
+## Usage
+- Just Click on `open in colab` button on top of this README file
+
+
+## Citation
+**Work in progress**

client-side/app.py

@@ -0,0 +1,38 @@
+from flask import Flask, request, jsonify, render_template
+from PIL import Image
+import requests
+from io import BytesIO
+import base64
+
+app = Flask(__name__)
+
+
+@app.route('/')
+def home():
+    return render_template("index.html")
+
+
+@app.route("/preds", methods=['POST'])
+def submit():
+    cloth = request.files['cloth']
+    model = request.files['model']
+
+    ## replace the url from the ngrok url provided on the notebook on server.
+    url = "http://e793-34-123-73-186.ngrok-free.app/api/transform"
+    print("sending")
+    response = requests.post(url=url, files={"cloth":cloth.stream, "model":model.stream})
+    op = Image.open(BytesIO(response.content))
+
+    buffer = BytesIO()
+    op.save(buffer, 'png')
+    buffer.seek(0)
+
+    data = buffer.read()
+    data = base64.b64encode(data).decode()
+
+
+    return render_template('index.html', op=data)
+    # return render_template('index.html', test=True)
+
+if __name__ == '__main__':
+    app.run(debug=True)

client-side/templates/index.html

@@ -0,0 +1,249 @@
+<!doctype html>
+<html>
+
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+
+    <script src="https://cdn.tailwindcss.com"></script>
+</head>
+
+<body>
+
+    <!-- <h1>Welcome to your own virtual clothing assistant</h1>
+    <form action="{{ url_for('submit') }}" method="post" enctype="multipart/form-data">
+        <input type="file" name="cloth">
+        <input type="file" name="model">
+        <button type="submit">Submit</button>
+    </form> -->
+
+
+
+    <header class="text-gray-400 bg-gray-900 body-font">
+        <div class="container mx-auto flex flex-wrap p-5 flex-col md:flex-row items-center">
+            <a class="flex title-font font-medium items-center text-white mb-4 md:mb-0">
+
+                <img src="{{url_for('static', filename='images/logo.png')}}" height="50" width="50" style="border-radius: 50%;"/>
+
+                <span class="ml-3 text-xl">Virtual Cloth Assistant</span>
+            </a>
+
+
+        </div>
+    </header>
+
+
+    <section class="text-gray-400 bg-gray-900 body-font">
+        <div class="container px-5 py-24 mx-auto">
+            <div class="text-center mb-20">
+                <h1 class="sm:text-3xl text-2xl font-medium title-font text-white mb-4">Virtual Cloth Assistant</h1>
+                <p class="text-base leading-relaxed xl:w-2/4 lg:w-3/4 mx-auto text-gray-400 text-opacity-80">Wanna try out, How that cloth suits you ?
+                    <br>
+                    Upgrade your shopping experience with an intelligent trial room.
+                    <br> Check out our API and get your wish fulfilled in seconds!!</p>
+                <div class="flex mt-6 justify-center">
+                    <div class="w-16 h-1 rounded-full bg-blue-500 inline-flex"></div>
+                </div>
+            </div>
+            <div class="flex flex-wrap sm:-m-4 -mx-4 -mb-10 -mt-4 md:space-y-0 space-y-6">
+                <div class="p-4 md:w-1/3 flex flex-col text-center items-center">
+                    <div
+                        class="w-20 h-20 inline-flex items-center justify-center rounded-full bg-gray-800 text-blue-400 mb-5 flex-shrink-0">
+                        <svg fill="none" stroke="currentColor" stroke-linecap="round" stroke-linejoin="round"
+                            stroke-width="2" class="w-10 h-10" viewBox="0 0 24 24">
+                            <path d="M22 12h-4l-3 9L9 3l-3 9H2"></path>
+                        </svg>
+                    </div>
+                    <div class="flex-grow">
+                        <h2 class="text-white text-lg title-font font-medium mb-3">The Problem</h2>
+                        <p class="leading-relaxed text-base">While buying clothes online, it is difficult for a customer to select a desirable outfit in the first attempt because they can’t try on clothes before they are delivered physically.
+
+                        </p>
+
+                    </div>
+                </div>
+                <div class="p-4 md:w-1/3 flex flex-col text-center items-center">
+                    <div
+                        class="w-20 h-20 inline-flex items-center justify-center rounded-full bg-gray-800 text-blue-400 mb-5 flex-shrink-0">
+                        <svg fill="none" stroke="currentColor" stroke-linecap="round" stroke-linejoin="round"
+                            stroke-width="2" class="w-10 h-10" viewBox="0 0 24 24">
+                            <circle cx="6" cy="6" r="3"></circle>
+                            <circle cx="6" cy="18" r="3"></circle>
+                            <path d="M20 4L8.12 15.88M14.47 14.48L20 20M8.12 8.12L12 12"></path>
+                        </svg>
+                    </div>
+                    <div class="flex-grow">
+                        <h2 class="text-white text-lg title-font font-medium mb-3">The Solution</h2>
+                        <p class="leading-relaxed text-base">E-commerce websites can be equipped with virtual trial rooms that allow users to try on multiple clothes virtually and select the best looking outfit in a single attempt.
+
+                        </p>
+
+                    </div>
+                </div>
+                <div class="p-4 md:w-1/3 flex flex-col text-center items-center">
+                    <div
+                        class="w-20 h-20 inline-flex items-center justify-center rounded-full bg-gray-800 text-blue-400 mb-5 flex-shrink-0">
+                        <svg fill="none" stroke="currentColor" stroke-linecap="round" stroke-linejoin="round"
+                            stroke-width="2" class="w-10 h-10" viewBox="0 0 24 24">
+                            <path d="M20 21v-2a4 4 0 00-4-4H8a4 4 0 00-4 4v2"></path>
+                            <circle cx="12" cy="7" r="4"></circle>
+                        </svg>
+                    </div>
+                    <div class="flex-grow">
+                        <h2 class="text-white text-lg title-font font-medium mb-3">The Approach</h2>
+                        <p class="leading-relaxed text-base">
+                            We used Deep Learning to solve this problem. VCA (virtual clothing assistant) for consumers, where
+                            user can select the cloth he/she wants to wear and then upload his/her image of any pose they want and VCA is capable to
+                            dress that human with his/her selected cloth.
+                        </p>
+
+                    </div>
+                </div>
+            </div>
+            <!-- <button
+                class="flex mx-auto mt-16 text-white bg-blue-500 border-0 py-2 px-8 focus:outline-none hover:bg-blue-600 rounded text-lg">Button</button> -->
+        </div>
+    </section>
+
+
+
+    <section class="text-gray-400 bg-gray-900 body-font">
+        <form action="{{ url_for('submit') }}" method="post" enctype="multipart/form-data">
+            <div class="container mx-auto flex flex-col px-5 py-24 justify-center items-center">
+
+                <div class="flex flex-wrap -m-2">
+                    <div class="p-1 xl:w-1/2 md:w-1/2 w-full">
+                        <center><label class="block text-lr font-medium text-white-700"> Cloth Image </label></center>
+                        <div
+                            class="mt-1 flex justify-center px-6 pt-5 pb-6 border-2 border-gray-300 border-dashed rounded-md">
+                            <div class="space-y-1 text-center">
+                                <svg class="mx-auto h-12 w-12 text-gray-400" stroke="currentColor" fill="none"
+                                    viewBox="0 0 48 48" aria-hidden="true">
+                                    <path
+                                        d="M28 8H12a4 4 0 00-4 4v20m32-12v8m0 0v8a4 4 0 01-4 4H12a4 4 0 01-4-4v-4m32-4l-3.172-3.172a4 4 0 00-5.656 0L28 28M8 32l9.172-9.172a4 4 0 015.656 0L28 28m0 0l4 4m4-24h8m-4-4v8m-12 4h.02"
+                                        stroke-width="2" stroke-linecap="round" stroke-linejoin="round" />
+                                </svg>
+                                <div class="flex text-sm text-gray-600">
+                                    <!-- <label for="file-upload"
+                                        class="relative cursor-pointer rounded-md font-medium text-indigo-600 hover:text-indigo-500 focus-within:outline-none focus-within:ring-2 focus-within:ring-offset-2 focus-within:ring-indigo-500">
+                                        <span>Upload a file</span> -->
+                                        <input class="block w-full text-sm text-slate-500
+                                        file:mr-4 file:py-2 file:px-4
+                                        file:rounded-full file:border-0
+                                        file:text-sm file:font-semibold
+                                        file:bg-violet-50 file:text-violet-700
+                                        hover:file:bg-violet-100" id="file-upload" type="file" name="cloth" class="sr-only">
+                                    </label>
+                                    <p class="pl-1">or drag and drop</p>
+                                </div>
+                                <p class="text-xs text-gray-500">PNG, JPG up to 10MB</p>
+                            </div>
+                        </div>
+
+                    </div>
+                    <div class="p-1 xl:w-1/2 md:w-1/2 w-full">
+                        <center> <label class="block text-lr font-medium text-white-700"> Model Image </label></center>
+                        <div
+                            class="mt-1 flex justify-center px-6 pt-5 pb-6 border-2 border-gray-300 border-dashed rounded-md">
+                            <div class="space-y-1 text-center">
+                                <svg class="mx-auto h-12 w-12 text-gray-400" stroke="currentColor" fill="none"
+                                    viewBox="0 0 48 48" aria-hidden="true">
+                                    <path
+                                        d="M28 8H12a4 4 0 00-4 4v20m32-12v8m0 0v8a4 4 0 01-4 4H12a4 4 0 01-4-4v-4m32-4l-3.172-3.172a4 4 0 00-5.656 0L28 28M8 32l9.172-9.172a4 4 0 015.656 0L28 28m0 0l4 4m4-24h8m-4-4v8m-12 4h.02"
+                                        stroke-width="2" stroke-linecap="round" stroke-linejoin="round" />
+                                </svg>
+                                <div class="flex text-sm text-gray-600">
+                                    <!-- <label for="file-upload"
+                                        class="relative cursor-pointer rounded-md font-medium text-indigo-600 hover:text-indigo-500 focus-within:outline-none focus-within:ring-2 focus-within:ring-offset-2 focus-within:ring-indigo-500">
+                                        <span>Upload a file</span> -->
+                                        <input class="block w-full text-sm text-slate-500
+                                        file:mr-4 file:py-2 file:px-4
+                                        file:rounded-full file:border-0
+                                        file:text-sm file:font-semibold
+                                        file:bg-violet-50 file:text-violet-700
+                                        hover:file:bg-violet-100" id="file-upload" type="file" name="model" class="sr-only">
+                                    </label>
+                                    <p class="pl-1">or drag and drop</p>
+                                </div>
+                                <p class="text-xs text-gray-500">PNG, JPG up to 10MB</p>
+                            </div>
+                        </div>
+                    </div>
+                </div>
+
+                <br>
+                <br>
+
+                <div class="w-full md:w-2/3 flex flex-col mb-16 items-center text-center">
+                    <h1 class="title-font sm:text-4xl text-3xl mb-4 font-medium text-white">Upload corresponding images
+                        and get the Result</h1>
+                    <p class="mb-8 leading-relaxed"> Need to wait for 20-30 seconds so chill and drink some water....
+                    </p>
+                    <div class="flex w-full justify-center items-end">
+
+                        <button type="submit"
+                            class="inline-flex text-white bg-blue-500 border-0 py-2 px-6 focus:outline-none hover:bg-blue-600 rounded text-lg">Try
+                            it</button>
+                    </div>
+
+                    </button>
+                </div>
+
+
+            </div>
+
+            <div>
+                {% if op %}
+                <center style="color: white; font-size: x-large;">HERE'S IS YOUR RESULT 🤗</center>
+
+                <center>
+                    <div class="sm: w-3/4 mb-10 lg:mb-0 rounded-lg overflow-hidden">
+                        <img alt="output" class="object-cover object-center h-2/4 w-2/4"
+                            src="data:image/png;base64,{{ op }}">
+                    </div>
+                </center>
+                {% endif %}
+            </div>
+        </form>
+
+
+    </section>
+
+
+    <footer class="text-gray-400 bg-gray-900 body-font">
+        <div
+            class="container px-5 py-24 mx-auto flex md:items-center lg:items-start md:flex-row md:flex-nowrap flex-wrap flex-col">
+            <div class="w-64 flex-shrink-0 md:mx-0 mx-auto text-center md:text-left md:mt-0 mt-10">
+                <a class="flex title-font font-medium items-center md:justify-start justify-center text-white">
+                    <img src="{{url_for('static', filename='images/logo.png')}}" height="50" width="50" style="border-radius: 50%;"/>
+                    <span class="ml-3 text-xl">V-Cloth Assistant</span>
+                </a>
+            </div>
+            <div class="flex-grow flex flex-wrap md:pr-20 -mb-10 md:text-left text-center order-first">
+                <div class="lg:w-1/4 md:w-1/2 w-full px-4">
+                    <h2 class="title-font font-medium text-white tracking-widest text-sm mb-3">SWAYAM</h2>
+                </div>
+                <div class="lg:w-1/4 md:w-1/2 w-full px-4">
+                    <h2 class="title-font font-medium text-white tracking-widest text-sm mb-3">PARTH</h2>
+                </div>
+                <div class="lg:w-1/4 md:w-1/2 w-full px-4">
+                    <h2 class="title-font font-medium text-white tracking-widest text-sm mb-3">KEERTHI</h2>
+                </div>
+                <div class="lg:w-1/4 md:w-1/2 w-full px-4">
+                    <h2 class="title-font font-medium text-white tracking-widest text-sm mb-3">NAVANEETH</h2>
+                </div>
+            </div>
+        </div>
+        <div class="bg-gray-800 bg-opacity-75">
+            <div class="container mx-auto py-4 px-5 flex flex-wrap flex-col sm:flex-row">
+                <p class="text-gray-400 text-sm text-center sm:text-left">© 2022 Crework Batch 3 —
+                    <a href="https://twitter.com/knyttneve" class="text-gray-500 ml-1" rel="noopener noreferrer"
+                        target="_blank">@Crework</a>
+                </p>
+            </div>
+        </div>
+    </footer>
+
+</body>
+
+</html>

cloth-mask.py

@@ -0,0 +1,124 @@
+import os
+from PIL import Image
+import numpy as np
+from collections import OrderedDict
+
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+
+from networks.u2net import U2NET
+device = 'cuda'
+
+image_dir = '/content/inputs/test/cloth'
+result_dir = '/content/inputs/test/cloth-mask'
+checkpoint_path = 'cloth_segm_u2net_latest.pth'
+
+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
+
+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"
+
+
+def get_palette(num_cls):
+    """ Returns the color map for visualizing the segmentation mask.
+    Args:
+        num_cls: Number of classes
+    Returns:
+        The color map
+    """
+    n = num_cls
+    palette = [0] * (n * 3)
+    for j in range(0, n):
+        lab = j
+        palette[j * 3 + 0] = 0
+        palette[j * 3 + 1] = 0
+        palette[j * 3 + 2] = 0
+        i = 0
+        while lab:
+            palette[j * 3 + 0] = 255
+            palette[j * 3 + 1] = 255
+            palette[j * 3 + 2] = 255
+            # palette[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i))
+            # palette[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i))
+            # palette[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i))
+            i += 1
+            lab >>= 3
+    return palette
+
+
+transforms_list = []
+transforms_list += [transforms.ToTensor()]
+transforms_list += [Normalize_image(0.5, 0.5)]
+transform_rgb = transforms.Compose(transforms_list)
+
+net = U2NET(in_ch=3, out_ch=4)
+net = load_checkpoint_mgpu(net, checkpoint_path)
+net = net.to(device)
+net = net.eval()
+
+palette = get_palette(4)
+
+images_list = sorted(os.listdir(image_dir))
+for image_name in images_list:
+    img = Image.open(os.path.join(image_dir, image_name)).convert('RGB')
+    img_size = img.size
+    img = img.resize((768, 768), Image.BICUBIC)
+    image_tensor = transform_rgb(img)
+    image_tensor = torch.unsqueeze(image_tensor, 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()
+
+    output_img = Image.fromarray(output_arr.astype('uint8'), mode='L')
+    output_img = output_img.resize(img_size, Image.BICUBIC)
+    
+    output_img.putpalette(palette)
+    output_img = output_img.convert('L')
+    output_img.save(os.path.join(result_dir, image_name[:-4]+'.jpg'))

datasets.py

@@ -0,0 +1,224 @@
+import json
+from os import path as osp
+
+import numpy as np
+from PIL import Image, ImageDraw
+import torch
+from torch.utils import data
+from torchvision import transforms
+
+
+class VITONDataset(data.Dataset):
+    def __init__(self, opt):
+        super(VITONDataset, self).__init__()
+        self.load_height = opt.load_height
+        self.load_width = opt.load_width
+        self.semantic_nc = opt.semantic_nc
+        self.data_path = osp.join(opt.dataset_dir, opt.dataset_mode)
+        self.transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+        ])
+
+        # load data list
+        img_names = []
+        c_names = []
+        with open(osp.join(opt.dataset_dir, opt.dataset_list), 'r') as f:
+            for line in f.readlines():
+                img_name, c_name = line.strip().split()
+                img_names.append(img_name)
+                c_names.append(c_name)
+
+        self.img_names = img_names
+        self.c_names = dict()
+        self.c_names['unpaired'] = c_names
+
+    def get_parse_agnostic(self, parse, pose_data):
+        parse_array = np.array(parse)
+        parse_upper = ((parse_array == 5).astype(np.float32) +
+                       (parse_array == 6).astype(np.float32) +
+                       (parse_array == 7).astype(np.float32))
+        parse_neck = (parse_array == 10).astype(np.float32)
+
+        r = 10
+        agnostic = parse.copy()
+
+        # mask arms
+        for parse_id, pose_ids in [(14, [2, 5, 6, 7]), (15, [5, 2, 3, 4])]:
+            mask_arm = Image.new('L', (self.load_width, self.load_height), 'black')
+            mask_arm_draw = ImageDraw.Draw(mask_arm)
+            i_prev = pose_ids[0]
+            for i in pose_ids[1:]:
+                if (pose_data[i_prev, 0] == 0.0 and pose_data[i_prev, 1] == 0.0) or (pose_data[i, 0] == 0.0 and pose_data[i, 1] == 0.0):
+                    continue
+                mask_arm_draw.line([tuple(pose_data[j]) for j in [i_prev, i]], 'white', width=r*10)
+                pointx, pointy = pose_data[i]
+                radius = r*4 if i == pose_ids[-1] else r*15
+                mask_arm_draw.ellipse((pointx-radius, pointy-radius, pointx+radius, pointy+radius), 'white', 'white')
+                i_prev = i
+            parse_arm = (np.array(mask_arm) / 255) * (parse_array == parse_id).astype(np.float32)
+            agnostic.paste(0, None, Image.fromarray(np.uint8(parse_arm * 255), 'L'))
+
+        # mask torso & neck
+        agnostic.paste(0, None, Image.fromarray(np.uint8(parse_upper * 255), 'L'))
+        agnostic.paste(0, None, Image.fromarray(np.uint8(parse_neck * 255), 'L'))
+
+        return agnostic
+
+    def get_img_agnostic(self, img, parse, pose_data):
+        parse_array = np.array(parse)
+        parse_head = ((parse_array == 4).astype(np.float32) +
+                      (parse_array == 13).astype(np.float32))
+        parse_lower = ((parse_array == 9).astype(np.float32) +
+                       (parse_array == 12).astype(np.float32) +
+                       (parse_array == 16).astype(np.float32) +
+                       (parse_array == 17).astype(np.float32) +
+                       (parse_array == 18).astype(np.float32) +
+                       (parse_array == 19).astype(np.float32))
+
+        r = 20
+        agnostic = img.copy()
+        agnostic_draw = ImageDraw.Draw(agnostic)
+
+        length_a = np.linalg.norm(pose_data[5] - pose_data[2])
+        length_b = np.linalg.norm(pose_data[12] - pose_data[9])
+        point = (pose_data[9] + pose_data[12]) / 2
+        pose_data[9] = point + (pose_data[9] - point) / length_b * length_a
+        pose_data[12] = point + (pose_data[12] - point) / length_b * length_a
+
+        # mask arms
+        agnostic_draw.line([tuple(pose_data[i]) for i in [2, 5]], 'gray', width=r*10)
+        for i in [2, 5]:
+            pointx, pointy = pose_data[i]
+            agnostic_draw.ellipse((pointx-r*5, pointy-r*5, pointx+r*5, pointy+r*5), 'gray', 'gray')
+        for i in [3, 4, 6, 7]:
+            if (pose_data[i - 1, 0] == 0.0 and pose_data[i - 1, 1] == 0.0) or (pose_data[i, 0] == 0.0 and pose_data[i, 1] == 0.0):
+                continue
+            agnostic_draw.line([tuple(pose_data[j]) for j in [i - 1, i]], 'gray', width=r*10)
+            pointx, pointy = pose_data[i]
+            agnostic_draw.ellipse((pointx-r*5, pointy-r*5, pointx+r*5, pointy+r*5), 'gray', 'gray')
+
+        # mask torso
+        for i in [9, 12]:
+            pointx, pointy = pose_data[i]
+            agnostic_draw.ellipse((pointx-r*3, pointy-r*6, pointx+r*3, pointy+r*6), 'gray', 'gray')
+        agnostic_draw.line([tuple(pose_data[i]) for i in [2, 9]], 'gray', width=r*6)
+        agnostic_draw.line([tuple(pose_data[i]) for i in [5, 12]], 'gray', width=r*6)
+        agnostic_draw.line([tuple(pose_data[i]) for i in [9, 12]], 'gray', width=r*12)
+        agnostic_draw.polygon([tuple(pose_data[i]) for i in [2, 5, 12, 9]], 'gray', 'gray')
+
+        # mask neck
+        pointx, pointy = pose_data[1]
+        agnostic_draw.rectangle((pointx-r*7, pointy-r*7, pointx+r*7, pointy+r*7), 'gray', 'gray')
+        agnostic.paste(img, None, Image.fromarray(np.uint8(parse_head * 255), 'L'))
+        agnostic.paste(img, None, Image.fromarray(np.uint8(parse_lower * 255), 'L'))
+
+        return agnostic
+
+    def __getitem__(self, index):
+        img_name = self.img_names[index]
+        c_name = {}
+        c = {}
+        cm = {}
+        for key in self.c_names:
+            c_name[key] = self.c_names[key][index]
+            c[key] = Image.open(osp.join(self.data_path, 'cloth', c_name[key])).convert('RGB')
+            c[key] = transforms.Resize(self.load_width, interpolation=2)(c[key])
+            cm[key] = Image.open(osp.join(self.data_path, 'cloth-mask', c_name[key]))
+            cm[key] = transforms.Resize(self.load_width, interpolation=0)(cm[key])
+
+            c[key] = self.transform(c[key])  # [-1,1]
+            cm_array = np.array(cm[key])
+            cm_array = (cm_array >= 128).astype(np.float32)
+            cm[key] = torch.from_numpy(cm_array)  # [0,1]
+            cm[key].unsqueeze_(0)
+
+        # load pose image
+        pose_name = img_name.replace('.jpg', '_rendered.png')
+        pose_rgb = Image.open(osp.join(self.data_path, 'openpose-img', pose_name))
+        pose_rgb = transforms.Resize(self.load_width, interpolation=2)(pose_rgb)
+        pose_rgb = self.transform(pose_rgb)  # [-1,1]
+
+        pose_name = img_name.replace('.jpg', '_keypoints.json')
+        with open(osp.join(self.data_path, 'openpose-json', pose_name), 'r') as f:
+            pose_label = json.load(f)
+            pose_data = pose_label['people'][0]['pose_keypoints_2d']
+            pose_data = np.array(pose_data)
+            pose_data = pose_data.reshape((-1, 3))[:, :2]
+
+        # load parsing image
+        parse_name = img_name.replace('.jpg', '.png')
+        parse = Image.open(osp.join(self.data_path, 'image-parse', parse_name))
+        parse = transforms.Resize(self.load_width, interpolation=0)(parse)
+        parse_agnostic = self.get_parse_agnostic(parse, pose_data)
+        parse_agnostic = torch.from_numpy(np.array(parse_agnostic)[None]).long()
+
+        labels = {
+            0: ['background', [0, 10]],
+            1: ['hair', [1, 2]],
+            2: ['face', [4, 13]],
+            3: ['upper', [5, 6, 7]],
+            4: ['bottom', [9, 12]],
+            5: ['left_arm', [14]],
+            6: ['right_arm', [15]],
+            7: ['left_leg', [16]],
+            8: ['right_leg', [17]],
+            9: ['left_shoe', [18]],
+            10: ['right_shoe', [19]],
+            11: ['socks', [8]],
+            12: ['noise', [3, 11]]
+        }
+        parse_agnostic_map = torch.zeros(20, self.load_height, self.load_width, dtype=torch.float)
+        parse_agnostic_map.scatter_(0, parse_agnostic, 1.0)
+        new_parse_agnostic_map = torch.zeros(self.semantic_nc, self.load_height, self.load_width, dtype=torch.float)
+        for i in range(len(labels)):
+            for label in labels[i][1]:
+                new_parse_agnostic_map[i] += parse_agnostic_map[label]
+
+        # load person image
+        img = Image.open(osp.join(self.data_path, 'image', img_name))
+        img = transforms.Resize(self.load_width, interpolation=2)(img)
+        img_agnostic = self.get_img_agnostic(img, parse, pose_data)
+        img = self.transform(img)
+        img_agnostic = self.transform(img_agnostic)  # [-1,1]
+
+        result = {
+            'img_name': img_name,
+            'c_name': c_name,
+            'img': img,
+            'img_agnostic': img_agnostic,
+            'parse_agnostic': new_parse_agnostic_map,
+            'pose': pose_rgb,
+            'cloth': c,
+            'cloth_mask': cm,
+        }
+        return result
+
+    def __len__(self):
+        return len(self.img_names)
+
+
+class VITONDataLoader:
+    def __init__(self, opt, dataset):
+        super(VITONDataLoader, self).__init__()
+
+        if opt.shuffle:
+            train_sampler = data.sampler.RandomSampler(dataset)
+        else:
+            train_sampler = None
+
+        self.data_loader = data.DataLoader(
+                dataset, batch_size=opt.batch_size, shuffle=(train_sampler is None),
+                num_workers=opt.workers, pin_memory=True, drop_last=True, sampler=train_sampler
+        )
+        self.dataset = dataset
+        self.data_iter = self.data_loader.__iter__()
+
+    def next_batch(self):
+        try:
+            batch = self.data_iter.__next__()
+        except StopIteration:
+            self.data_iter = self.data_loader.__iter__()
+            batch = self.data_iter.__next__()
+
+        return batch

network.py

@@ -0,0 +1,526 @@
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn import init
+from torch.nn.utils.spectral_norm import spectral_norm
+
+
+# ----------------------------------------------------------------------------------------------------------------------
+#                                                     Common classes
+# ----------------------------------------------------------------------------------------------------------------------
+class BaseNetwork(nn.Module):
+    def __init__(self):
+        super(BaseNetwork, self).__init__()
+
+    def print_network(self):
+        num_params = 0
+        for param in self.parameters():
+            num_params += param.numel()
+        print("Network [{}] was created. Total number of parameters: {:.1f} million. "
+              "To see the architecture, do print(network).".format(self.__class__.__name__, num_params / 1000000))
+
+    def init_weights(self, init_type='normal', gain=0.02):
+        def init_func(m):
+            classname = m.__class__.__name__
+            if 'BatchNorm2d' in classname:
+                if hasattr(m, 'weight') and m.weight is not None:
+                    init.normal_(m.weight.data, 1.0, gain)
+                if hasattr(m, 'bias') and m.bias is not None:
+                    init.constant_(m.bias.data, 0.0)
+            elif ('Conv' in classname or 'Linear' in classname) and hasattr(m, 'weight'):
+                if init_type == 'normal':
+                    init.normal_(m.weight.data, 0.0, gain)
+                elif init_type == 'xavier':
+                    init.xavier_normal_(m.weight.data, gain=gain)
+                elif init_type == 'xavier_uniform':
+                    init.xavier_uniform_(m.weight.data, gain=1.0)
+                elif init_type == 'kaiming':
+                    init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+                elif init_type == 'orthogonal':
+                    init.orthogonal_(m.weight.data, gain=gain)
+                elif init_type == 'none':  # uses pytorch's default init method
+                    m.reset_parameters()
+                else:
+                    raise NotImplementedError("initialization method '{}' is not implemented".format(init_type))
+                if hasattr(m, 'bias') and m.bias is not None:
+                    init.constant_(m.bias.data, 0.0)
+
+        self.apply(init_func)
+
+    def forward(self, *inputs):
+        pass
+
+
+# ----------------------------------------------------------------------------------------------------------------------
+#                                              SegGenerator-related classes
+# ----------------------------------------------------------------------------------------------------------------------
+class SegGenerator(BaseNetwork):
+    def __init__(self, opt, input_nc, output_nc=13, norm_layer=nn.InstanceNorm2d):
+        super(SegGenerator, self).__init__()
+
+        self.conv1 = nn.Sequential(nn.Conv2d(input_nc, 64, kernel_size=3, padding=1), norm_layer(64), nn.ReLU(),
+                                   nn.Conv2d(64, 64, kernel_size=3, padding=1), norm_layer(64), nn.ReLU())
+
+        self.conv2 = nn.Sequential(nn.Conv2d(64, 128, kernel_size=3, padding=1), norm_layer(128), nn.ReLU(),
+                                   nn.Conv2d(128, 128, kernel_size=3, padding=1), norm_layer(128), nn.ReLU())
+
+        self.conv3 = nn.Sequential(nn.Conv2d(128, 256, kernel_size=3, padding=1), norm_layer(256), nn.ReLU(),
+                                   nn.Conv2d(256, 256, kernel_size=3, padding=1), norm_layer(256), nn.ReLU())
+
+        self.conv4 = nn.Sequential(nn.Conv2d(256, 512, kernel_size=3, padding=1), norm_layer(512), nn.ReLU(),
+                                   nn.Conv2d(512, 512, kernel_size=3, padding=1), norm_layer(512), nn.ReLU())
+
+        self.conv5 = nn.Sequential(nn.Conv2d(512, 1024, kernel_size=3, padding=1), norm_layer(1024), nn.ReLU(),
+                                   nn.Conv2d(1024, 1024, kernel_size=3, padding=1), norm_layer(1024), nn.ReLU())
+
+        self.up6 = nn.Sequential(nn.Upsample(scale_factor=2, mode='nearest'),
+                                 nn.Conv2d(1024, 512, kernel_size=3, padding=1), norm_layer(512), nn.ReLU())
+        self.conv6 = nn.Sequential(nn.Conv2d(1024, 512, kernel_size=3, padding=1), norm_layer(512), nn.ReLU(),
+                                   nn.Conv2d(512, 512, kernel_size=3, padding=1), norm_layer(512), nn.ReLU())
+
+        self.up7 = nn.Sequential(nn.Upsample(scale_factor=2, mode='nearest'),
+                                 nn.Conv2d(512, 256, kernel_size=3, padding=1), norm_layer(256), nn.ReLU())
+        self.conv7 = nn.Sequential(nn.Conv2d(512, 256, kernel_size=3, padding=1), norm_layer(256), nn.ReLU(),
+                                   nn.Conv2d(256, 256, kernel_size=3, padding=1), norm_layer(256), nn.ReLU())
+
+        self.up8 = nn.Sequential(nn.Upsample(scale_factor=2, mode='nearest'),
+                                 nn.Conv2d(256, 128, kernel_size=3, padding=1), norm_layer(128), nn.ReLU())
+        self.conv8 = nn.Sequential(nn.Conv2d(256, 128, kernel_size=3, padding=1), norm_layer(128), nn.ReLU(),
+                                   nn.Conv2d(128, 128, kernel_size=3, padding=1), norm_layer(128), nn.ReLU())
+
+        self.up9 = nn.Sequential(nn.Upsample(scale_factor=2, mode='nearest'),
+                                 nn.Conv2d(128, 64, kernel_size=3, padding=1), norm_layer(64), nn.ReLU())
+        self.conv9 = nn.Sequential(nn.Conv2d(128, 64, kernel_size=3, padding=1), norm_layer(64), nn.ReLU(),
+                                   nn.Conv2d(64, 64, kernel_size=3, padding=1), norm_layer(64), nn.ReLU(),
+                                   nn.Conv2d(64, output_nc, kernel_size=3, padding=1))
+
+        self.pool = nn.MaxPool2d(2)
+        self.drop = nn.Dropout(0.5)
+        self.sigmoid = nn.Sigmoid()
+
+        self.print_network()
+        self.init_weights(opt.init_type, opt.init_variance)
+
+    def forward(self, x):
+        conv1 = self.conv1(x)
+        conv2 = self.conv2(self.pool(conv1))
+        conv3 = self.conv3(self.pool(conv2))
+        conv4 = self.drop(self.conv4(self.pool(conv3)))
+        conv5 = self.drop(self.conv5(self.pool(conv4)))
+
+        conv6 = self.conv6(torch.cat((conv4, self.up6(conv5)), 1))
+        conv7 = self.conv7(torch.cat((conv3, self.up7(conv6)), 1))
+        conv8 = self.conv8(torch.cat((conv2, self.up8(conv7)), 1))
+        conv9 = self.conv9(torch.cat((conv1, self.up9(conv8)), 1))
+        return self.sigmoid(conv9)
+
+
+# ----------------------------------------------------------------------------------------------------------------------
+#                                                  GMM-related classes
+# ----------------------------------------------------------------------------------------------------------------------
+class FeatureExtraction(BaseNetwork):
+    def __init__(self, input_nc, ngf=64, num_layers=4, norm_layer=nn.BatchNorm2d):
+        super(FeatureExtraction, self).__init__()
+
+        nf = ngf
+        layers = [nn.Conv2d(input_nc, nf, kernel_size=4, stride=2, padding=1), nn.ReLU(), norm_layer(nf)]
+
+        for i in range(1, num_layers):
+            nf_prev = nf
+            nf = min(nf * 2, 512)
+            layers += [nn.Conv2d(nf_prev, nf, kernel_size=4, stride=2, padding=1), nn.ReLU(), norm_layer(nf)]
+
+        layers += [nn.Conv2d(nf, 512, kernel_size=3, stride=1, padding=1), nn.ReLU(), norm_layer(512)]
+        layers += [nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1), nn.ReLU()]
+
+        self.model = nn.Sequential(*layers)
+        self.init_weights()
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class FeatureCorrelation(nn.Module):
+    def __init__(self):
+        super(FeatureCorrelation, self).__init__()
+
+    def forward(self, featureA, featureB):
+        # Reshape features for matrix multiplication.
+        b, c, h, w = featureA.size()
+        featureA = featureA.permute(0, 3, 2, 1).reshape(b, w * h, c)
+        featureB = featureB.reshape(b, c, h * w)
+
+        # Perform matrix multiplication.
+        corr = torch.bmm(featureA, featureB).reshape(b, w * h, h, w)
+        return corr
+
+
+class FeatureRegression(nn.Module):
+    def __init__(self, input_nc=512, output_size=6, norm_layer=nn.BatchNorm2d):
+        super(FeatureRegression, self).__init__()
+
+        self.conv = nn.Sequential(
+            nn.Conv2d(input_nc, 512, kernel_size=4, stride=2, padding=1), norm_layer(512), nn.ReLU(),
+            nn.Conv2d(512, 256, kernel_size=4, stride=2, padding=1), norm_layer(256), nn.ReLU(),
+            nn.Conv2d(256, 128, kernel_size=3, padding=1), norm_layer(128), nn.ReLU(),
+            nn.Conv2d(128, 64, kernel_size=3, padding=1), norm_layer(64), nn.ReLU()
+        )
+        self.linear = nn.Linear(64 * (input_nc // 16), output_size)
+        self.tanh = nn.Tanh()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.linear(x.reshape(x.size(0), -1))
+        return self.tanh(x)
+
+
+class TpsGridGen(nn.Module):
+    def __init__(self, opt, dtype=torch.float):
+        super(TpsGridGen, self).__init__()
+
+        # Create a grid in numpy.
+        # TODO: set an appropriate interval ([-1, 1] in CP-VTON, [-0.9, 0.9] in the current version of VITON-HD)
+        grid_X, grid_Y = np.meshgrid(np.linspace(-0.9, 0.9, opt.load_width), np.linspace(-0.9, 0.9, opt.load_height))
+        grid_X = torch.tensor(grid_X, dtype=dtype).unsqueeze(0).unsqueeze(3)  # size: (1, h, w, 1)
+        grid_Y = torch.tensor(grid_Y, dtype=dtype).unsqueeze(0).unsqueeze(3)  # size: (1, h, w, 1)
+
+        # Initialize the regular grid for control points P.
+        self.N = opt.grid_size * opt.grid_size
+        coords = np.linspace(-0.9, 0.9, opt.grid_size)
+        # FIXME: why P_Y and P_X are swapped?
+        P_Y, P_X = np.meshgrid(coords, coords)
+        P_X = torch.tensor(P_X, dtype=dtype).reshape(self.N, 1)
+        P_Y = torch.tensor(P_Y, dtype=dtype).reshape(self.N, 1)
+        P_X_base = P_X.clone()
+        P_Y_base = P_Y.clone()
+
+        Li = self.compute_L_inverse(P_X, P_Y).unsqueeze(0)
+        P_X = P_X.unsqueeze(2).unsqueeze(3).unsqueeze(4).transpose(0, 4)  # size: (1, 1, 1, 1, self.N)
+        P_Y = P_Y.unsqueeze(2).unsqueeze(3).unsqueeze(4).transpose(0, 4)  # size: (1, 1, 1, 1, self.N)
+
+        self.register_buffer('grid_X', grid_X, False)
+        self.register_buffer('grid_Y', grid_Y, False)
+        self.register_buffer('P_X_base', P_X_base, False)
+        self.register_buffer('P_Y_base', P_Y_base, False)
+        self.register_buffer('Li', Li, False)
+        self.register_buffer('P_X', P_X, False)
+        self.register_buffer('P_Y', P_Y, False)
+
+    # TODO: refactor
+    def compute_L_inverse(self,X,Y):
+        N = X.size()[0] # num of points (along dim 0)
+        # construct matrix K
+        Xmat = X.expand(N,N)
+        Ymat = Y.expand(N,N)
+        P_dist_squared = torch.pow(Xmat-Xmat.transpose(0,1),2)+torch.pow(Ymat-Ymat.transpose(0,1),2)
+        P_dist_squared[P_dist_squared==0]=1 # make diagonal 1 to avoid NaN in log computation
+        K = torch.mul(P_dist_squared,torch.log(P_dist_squared))
+        # construct matrix L
+        O = torch.FloatTensor(N,1).fill_(1)
+        Z = torch.FloatTensor(3,3).fill_(0)
+        P = torch.cat((O,X,Y),1)
+        L = torch.cat((torch.cat((K,P),1),torch.cat((P.transpose(0,1),Z),1)),0)
+        Li = torch.inverse(L)
+        return Li
+
+    # TODO: refactor
+    def apply_transformation(self,theta,points):
+        if theta.dim()==2:
+            theta = theta.unsqueeze(2).unsqueeze(3)
+        # points should be in the [B,H,W,2] format,
+        # where points[:,:,:,0] are the X coords
+        # and points[:,:,:,1] are the Y coords
+
+        # input are the corresponding control points P_i
+        batch_size = theta.size()[0]
+        # split theta into point coordinates
+        Q_X=theta[:,:self.N,:,:].squeeze(3)
+        Q_Y=theta[:,self.N:,:,:].squeeze(3)
+        Q_X = Q_X + self.P_X_base.expand_as(Q_X)
+        Q_Y = Q_Y + self.P_Y_base.expand_as(Q_Y)
+
+        # get spatial dimensions of points
+        points_b = points.size()[0]
+        points_h = points.size()[1]
+        points_w = points.size()[2]
+
+        # repeat pre-defined control points along spatial dimensions of points to be transformed
+        P_X = self.P_X.expand((1,points_h,points_w,1,self.N))
+        P_Y = self.P_Y.expand((1,points_h,points_w,1,self.N))
+
+        # compute weigths for non-linear part
+        W_X = torch.bmm(self.Li[:,:self.N,:self.N].expand((batch_size,self.N,self.N)),Q_X)
+        W_Y = torch.bmm(self.Li[:,:self.N,:self.N].expand((batch_size,self.N,self.N)),Q_Y)
+        # reshape
+        # W_X,W,Y: size [B,H,W,1,N]
+        W_X = W_X.unsqueeze(3).unsqueeze(4).transpose(1,4).repeat(1,points_h,points_w,1,1)
+        W_Y = W_Y.unsqueeze(3).unsqueeze(4).transpose(1,4).repeat(1,points_h,points_w,1,1)
+        # compute weights for affine part
+        A_X = torch.bmm(self.Li[:,self.N:,:self.N].expand((batch_size,3,self.N)),Q_X)
+        A_Y = torch.bmm(self.Li[:,self.N:,:self.N].expand((batch_size,3,self.N)),Q_Y)
+        # reshape
+        # A_X,A,Y: size [B,H,W,1,3]
+        A_X = A_X.unsqueeze(3).unsqueeze(4).transpose(1,4).repeat(1,points_h,points_w,1,1)
+        A_Y = A_Y.unsqueeze(3).unsqueeze(4).transpose(1,4).repeat(1,points_h,points_w,1,1)
+
+        # compute distance P_i - (grid_X,grid_Y)
+        # grid is expanded in point dim 4, but not in batch dim 0, as points P_X,P_Y are fixed for all batch
+        points_X_for_summation = points[:,:,:,0].unsqueeze(3).unsqueeze(4).expand(points[:,:,:,0].size()+(1,self.N))
+        points_Y_for_summation = points[:,:,:,1].unsqueeze(3).unsqueeze(4).expand(points[:,:,:,1].size()+(1,self.N))
+
+        if points_b==1:
+            delta_X = points_X_for_summation-P_X
+            delta_Y = points_Y_for_summation-P_Y
+        else:
+            # use expanded P_X,P_Y in batch dimension
+            delta_X = points_X_for_summation-P_X.expand_as(points_X_for_summation)
+            delta_Y = points_Y_for_summation-P_Y.expand_as(points_Y_for_summation)
+
+        dist_squared = torch.pow(delta_X,2)+torch.pow(delta_Y,2)
+        # U: size [1,H,W,1,N]
+        dist_squared[dist_squared==0]=1 # avoid NaN in log computation
+        U = torch.mul(dist_squared,torch.log(dist_squared))
+
+        # expand grid in batch dimension if necessary
+        points_X_batch = points[:,:,:,0].unsqueeze(3)
+        points_Y_batch = points[:,:,:,1].unsqueeze(3)
+        if points_b==1:
+            points_X_batch = points_X_batch.expand((batch_size,)+points_X_batch.size()[1:])
+            points_Y_batch = points_Y_batch.expand((batch_size,)+points_Y_batch.size()[1:])
+
+        points_X_prime = A_X[:,:,:,:,0]+ \
+                       torch.mul(A_X[:,:,:,:,1],points_X_batch) + \
+                       torch.mul(A_X[:,:,:,:,2],points_Y_batch) + \
+                       torch.sum(torch.mul(W_X,U.expand_as(W_X)),4)
+
+        points_Y_prime = A_Y[:,:,:,:,0]+ \
+                       torch.mul(A_Y[:,:,:,:,1],points_X_batch) + \
+                       torch.mul(A_Y[:,:,:,:,2],points_Y_batch) + \
+                       torch.sum(torch.mul(W_Y,U.expand_as(W_Y)),4)
+
+        return torch.cat((points_X_prime,points_Y_prime),3)
+
+    def forward(self, theta):
+        warped_grid = self.apply_transformation(theta, torch.cat((self.grid_X, self.grid_Y), 3))
+        return warped_grid
+
+
+class GMM(nn.Module):
+    def __init__(self, opt, inputA_nc, inputB_nc):
+        super(GMM, self).__init__()
+
+        self.extractionA = FeatureExtraction(inputA_nc, ngf=64, num_layers=4)
+        self.extractionB = FeatureExtraction(inputB_nc, ngf=64, num_layers=4)
+        self.correlation = FeatureCorrelation()
+        self.regression = FeatureRegression(input_nc=(opt.load_width // 64) * (opt.load_height // 64),
+                                            output_size=2 * opt.grid_size**2)
+        self.gridGen = TpsGridGen(opt)
+
+    def forward(self, inputA, inputB):
+        featureA = F.normalize(self.extractionA(inputA), dim=1)
+        featureB = F.normalize(self.extractionB(inputB), dim=1)
+        corr = self.correlation(featureA, featureB)
+        theta = self.regression(corr)
+
+        warped_grid = self.gridGen(theta)
+        return theta, warped_grid
+
+
+# ----------------------------------------------------------------------------------------------------------------------
+#                                             ALIASGenerator-related classes
+# ----------------------------------------------------------------------------------------------------------------------
+class MaskNorm(nn.Module):
+    def __init__(self, norm_nc):
+        super(MaskNorm, self).__init__()
+
+        self.norm_layer = nn.InstanceNorm2d(norm_nc, affine=False)
+
+    def normalize_region(self, region, mask):
+        b, c, h, w = region.size()
+
+        num_pixels = mask.sum((2, 3), keepdim=True)  # size: (b, 1, 1, 1)
+        num_pixels[num_pixels == 0] = 1
+        mu = region.sum((2, 3), keepdim=True) / num_pixels  # size: (b, c, 1, 1)
+
+        normalized_region = self.norm_layer(region + (1 - mask) * mu)
+        return normalized_region * torch.sqrt(num_pixels / (h * w))
+
+    def forward(self, x, mask):
+        mask = mask.detach()
+        normalized_foreground = self.normalize_region(x * mask, mask)
+        normalized_background = self.normalize_region(x * (1 - mask), 1 - mask)
+        return normalized_foreground + normalized_background
+
+
+class ALIASNorm(nn.Module):
+    def __init__(self, norm_type, norm_nc, label_nc):
+        super(ALIASNorm, self).__init__()
+
+        self.noise_scale = nn.Parameter(torch.zeros(norm_nc))
+
+        assert norm_type.startswith('alias')
+        param_free_norm_type = norm_type[len('alias'):]
+        if param_free_norm_type == 'batch':
+            self.param_free_norm = nn.BatchNorm2d(norm_nc, affine=False)
+        elif param_free_norm_type == 'instance':
+            self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)
+        elif param_free_norm_type == 'mask':
+            self.param_free_norm = MaskNorm(norm_nc)
+        else:
+            raise ValueError(
+                "'{}' is not a recognized parameter-free normalization type in ALIASNorm".format(param_free_norm_type)
+            )
+
+        nhidden = 128
+        ks = 3
+        pw = ks // 2
+        self.conv_shared = nn.Sequential(nn.Conv2d(label_nc, nhidden, kernel_size=ks, padding=pw), nn.ReLU())
+        self.conv_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw)
+        self.conv_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw)
+
+    def forward(self, x, seg, misalign_mask=None):
+        # Part 1. Generate parameter-free normalized activations.
+        b, c, h, w = x.size()
+        noise = (torch.randn(b, w, h, 1).cuda() * self.noise_scale).transpose(1, 3)
+
+        if misalign_mask is None:
+            normalized = self.param_free_norm(x + noise)
+        else:
+            normalized = self.param_free_norm(x + noise, misalign_mask)
+
+        # Part 2. Produce affine parameters conditioned on the segmentation map.
+        actv = self.conv_shared(seg)
+        gamma = self.conv_gamma(actv)
+        beta = self.conv_beta(actv)
+
+        # Apply the affine parameters.
+        output = normalized * (1 + gamma) + beta
+        return output
+
+
+class ALIASResBlock(nn.Module):
+    def __init__(self, opt, input_nc, output_nc, use_mask_norm=True):
+        super(ALIASResBlock, self).__init__()
+
+        self.learned_shortcut = (input_nc != output_nc)
+        middle_nc = min(input_nc, output_nc)
+
+        self.conv_0 = nn.Conv2d(input_nc, middle_nc, kernel_size=3, padding=1)
+        self.conv_1 = nn.Conv2d(middle_nc, output_nc, kernel_size=3, padding=1)
+        if self.learned_shortcut:
+            self.conv_s = nn.Conv2d(input_nc, output_nc, kernel_size=1, bias=False)
+
+        subnorm_type = opt.norm_G
+        if subnorm_type.startswith('spectral'):
+            subnorm_type = subnorm_type[len('spectral'):]
+            self.conv_0 = spectral_norm(self.conv_0)
+            self.conv_1 = spectral_norm(self.conv_1)
+            if self.learned_shortcut:
+                self.conv_s = spectral_norm(self.conv_s)
+
+        semantic_nc = opt.semantic_nc
+        if use_mask_norm:
+            subnorm_type = 'aliasmask'
+            semantic_nc = semantic_nc + 1
+
+        self.norm_0 = ALIASNorm(subnorm_type, input_nc, semantic_nc)
+        self.norm_1 = ALIASNorm(subnorm_type, middle_nc, semantic_nc)
+        if self.learned_shortcut:
+            self.norm_s = ALIASNorm(subnorm_type, input_nc, semantic_nc)
+
+        self.relu = nn.LeakyReLU(0.2)
+
+    def shortcut(self, x, seg, misalign_mask):
+        if self.learned_shortcut:
+            return self.conv_s(self.norm_s(x, seg, misalign_mask))
+        else:
+            return x
+
+    def forward(self, x, seg, misalign_mask=None):
+        seg = F.interpolate(seg, size=x.size()[2:], mode='nearest')
+        if misalign_mask is not None:
+            misalign_mask = F.interpolate(misalign_mask, size=x.size()[2:], mode='nearest')
+
+        x_s = self.shortcut(x, seg, misalign_mask)
+
+        dx = self.conv_0(self.relu(self.norm_0(x, seg, misalign_mask)))
+        dx = self.conv_1(self.relu(self.norm_1(dx, seg, misalign_mask)))
+        output = x_s + dx
+        return output
+
+
+class ALIASGenerator(BaseNetwork):
+    def __init__(self, opt, input_nc):
+        super(ALIASGenerator, self).__init__()
+        self.num_upsampling_layers = opt.num_upsampling_layers
+
+        self.sh, self.sw = self.compute_latent_vector_size(opt)
+
+        nf = opt.ngf
+        self.conv_0 = nn.Conv2d(input_nc, nf * 16, kernel_size=3, padding=1)
+        for i in range(1, 8):
+            self.add_module('conv_{}'.format(i), nn.Conv2d(input_nc, 16, kernel_size=3, padding=1))
+
+        self.head_0 = ALIASResBlock(opt, nf * 16, nf * 16)
+
+        self.G_middle_0 = ALIASResBlock(opt, nf * 16 + 16, nf * 16)
+        self.G_middle_1 = ALIASResBlock(opt, nf * 16 + 16, nf * 16)
+
+        self.up_0 = ALIASResBlock(opt, nf * 16 + 16, nf * 8)
+        self.up_1 = ALIASResBlock(opt, nf * 8 + 16, nf * 4)
+        self.up_2 = ALIASResBlock(opt, nf * 4 + 16, nf * 2, use_mask_norm=False)
+        self.up_3 = ALIASResBlock(opt, nf * 2 + 16, nf * 1, use_mask_norm=False)
+        if self.num_upsampling_layers == 'most':
+            self.up_4 = ALIASResBlock(opt, nf * 1 + 16, nf // 2, use_mask_norm=False)
+            nf = nf // 2
+
+        self.conv_img = nn.Conv2d(nf, 3, kernel_size=3, padding=1)
+
+        self.up = nn.Upsample(scale_factor=2, mode='nearest')
+        self.relu = nn.LeakyReLU(0.2)
+        self.tanh = nn.Tanh()
+
+        self.print_network()
+        self.init_weights(opt.init_type, opt.init_variance)
+
+    def compute_latent_vector_size(self, opt):
+        if self.num_upsampling_layers == 'normal':
+            num_up_layers = 5
+        elif self.num_upsampling_layers == 'more':
+            num_up_layers = 6
+        elif self.num_upsampling_layers == 'most':
+            num_up_layers = 7
+        else:
+            raise ValueError("opt.num_upsampling_layers '{}' is not recognized".format(self.num_upsampling_layers))
+
+        sh = opt.load_height // 2**num_up_layers
+        sw = opt.load_width // 2**num_up_layers
+        return sh, sw
+
+    def forward(self, x, seg, seg_div, misalign_mask):
+        samples = [F.interpolate(x, size=(self.sh * 2**i, self.sw * 2**i), mode='nearest') for i in range(8)]
+        features = [self._modules['conv_{}'.format(i)](samples[i]) for i in range(8)]
+
+        x = self.head_0(features[0], seg_div, misalign_mask)
+
+        x = self.up(x)
+        x = self.G_middle_0(torch.cat((x, features[1]), 1), seg_div, misalign_mask)
+        if self.num_upsampling_layers in ['more', 'most']:
+            x = self.up(x)
+        x = self.G_middle_1(torch.cat((x, features[2]), 1), seg_div, misalign_mask)
+
+        x = self.up(x)
+        x = self.up_0(torch.cat((x, features[3]), 1), seg_div, misalign_mask)
+        x = self.up(x)
+        x = self.up_1(torch.cat((x, features[4]), 1), seg_div, misalign_mask)
+        x = self.up(x)
+        x = self.up_2(torch.cat((x, features[5]), 1), seg)
+        x = self.up(x)
+        x = self.up_3(torch.cat((x, features[6]), 1), seg)
+        if self.num_upsampling_layers == 'most':
+            x = self.up(x)
+            x = self.up_4(torch.cat((x, features[7]), 1), seg)
+
+        x = self.conv_img(self.relu(x))
+        return self.tanh(x)

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

remove_bg.py

@@ -0,0 +1,59 @@
+import requests
+import os
+from PIL import Image
+import numpy as np
+from rembg import remove
+
+
+class preprcessInput:
+
+    def __init__(self):
+        self.o_width = None
+        self.o_height = None
+        self.o_image = None
+
+        self.t_width = None
+        self.t_height = None
+        self.t_image = None
+        self.save_path = None
+
+    def remove_bg(self, file_path: str):
+        self.save_path = file_path[:-3]+'.png'
+        pic = Image.open(file_path)
+        self.o_width = np.asarray(pic).shape[1]
+        self.o_height = np.asarray(pic).shape[0]
+        try:
+            self.o_channels = np.asarray(pic).shape[2]
+        except Exception as e:
+            print("Single channel image and error", e)
+        os.remove(file_path)
+        self.o_image = remove(pic)
+        self.o_image.save(self.save_path)
+        os.remove(self.save_path)
+        return np.asarray(self.o_image)
+
+    def transform(self, width=768, height=1024):
+        newsize = (width, height)
+        self.t_height = height
+        self.t_width = width
+
+        pic = self.o_image
+        img = pic.resize(newsize)
+
+        self.t_image = img
+
+        background = Image.new("RGBA", newsize, (255, 255, 255, 255))
+        background.paste(img, mask=img.split()[3])  # 3 is the alpha channel
+        self.save_path = self.save_path[:-3] + '.jpg'
+        background.convert('RGB').save(self.save_path, 'JPEG')
+
+        return np.asarray(background.convert('RGB'))
+
+
+# USAGE OF THE CLASS
+preprocess = preprcessInput()
+for images in os.listdir('/content/inputs/test/image'):
+    print(images)
+    if images[-3:] == 'jpg':
+        op = preprocess.remove_bg(r'/content/inputs/test/image/'+images)
+        arr = preprocess.transform(768, 1024)

run.py

@@ -0,0 +1,41 @@
+from PIL import Image
+import os
+
+
+# running the preprocessing
+
+def resize_img(path):
+    im = Image.open(path)
+    im = im.resize((768, 1024))
+    im.save(path)
+
+
+for path in os.listdir('/content/inputs/test/cloth/'):
+    resize_img(f'/content/inputs/test/cloth/{path}')
+
+os.chdir('/content/clothes-virtual-try-on')
+os.system("rm -rf /content/inputs/test/cloth/.ipynb_checkpoints")
+os.system("python cloth-mask.py")
+os.chdir('/content')
+os.system("python /content/clothes-virtual-try-on/remove_bg.py")
+os.system(
+    "python3 /content/Self-Correction-Human-Parsing/simple_extractor.py --dataset 'lip' --model-restore '/content/Self-Correction-Human-Parsing/checkpoints/final.pth' --input-dir '/content/inputs/test/image' --output-dir '/content/inputs/test/image-parse'")
+os.chdir('/content')
+os.system(
+    "cd openpose && ./build/examples/openpose/openpose.bin --image_dir /content/inputs/test/image/ --write_json /content/inputs/test/openpose-json/ --display 0 --render_pose 0 --hand")
+os.system(
+    "cd openpose && ./build/examples/openpose/openpose.bin --image_dir /content/inputs/test/image/ --display 0 --write_images /content/inputs/test/openpose-img/ --hand --render_pose 1 --disable_blending true")
+
+model_image = os.listdir('/content/inputs/test/image')
+cloth_image = os.listdir('/content/inputs/test/cloth')
+pairs = zip(model_image, cloth_image)
+
+with open('/content/inputs/test_pairs.txt', 'w') as file:
+    for model, cloth in pairs:
+        file.write(f"{model} {cloth}")
+
+# making predictions
+os.system(
+    "python /content/clothes-virtual-try-on/test.py --name output --dataset_dir /content/inputs --checkpoint_dir /content/clothes-virtual-try-on/checkpoints --save_dir /content/")
+os.system("rm -rf /content/inputs")
+os.system("rm -rf /content/output/.ipynb_checkpoints")

setup_ngrok.ipynb

@@ -0,0 +1,643 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8gqt11Y_RYxU"
+      },
+      "source": [
+        "# Setting up the environment. PLEASE WAIT 🙃"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "RHmGnnTZL6os"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install --upgrade --no-cache-dir gdown\n",
+        "!pip install rembg[gpu]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "uX3LsFFPKSwo"
+      },
+      "outputs": [],
+      "source": [
+        "! wget -c \"https://github.com/Kitware/CMake/releases/download/v3.19.6/cmake-3.19.6.tar.gz\"\n",
+        "! tar xf cmake-3.19.6.tar.gz\n",
+        "! cd cmake-3.19.6 && ./configure && make && sudo make install"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "51QJAhPOK9cK"
+      },
+      "outputs": [],
+      "source": [
+        "# Install library\n",
+        "! sudo apt-get --assume-yes update\n",
+        "! sudo apt-get --assume-yes install build-essential\n",
+        "# OpenCV\n",
+        "! sudo apt-get --assume-yes install libopencv-dev\n",
+        "# General dependencies\n",
+        "! sudo apt-get --assume-yes install libatlas-base-dev libprotobuf-dev libleveldb-dev libsnappy-dev libhdf5-serial-dev protobuf-compiler\n",
+        "! sudo apt-get --assume-yes install --no-install-recommends libboost-all-dev\n",
+        "# Remaining dependencies, 14.04\n",
+        "! sudo apt-get --assume-yes install libgflags-dev libgoogle-glog-dev liblmdb-dev\n",
+        "# Python3 libs\n",
+        "! sudo apt-get --assume-yes install python3-setuptools python3-dev build-essential\n",
+        "! sudo apt-get --assume-yes install python3-pip\n",
+        "! sudo -H pip3 install --upgrade numpy protobuf opencv-python\n",
+        "# OpenCL Generic\n",
+        "! sudo apt-get --assume-yes install opencl-headers ocl-icd-opencl-dev\n",
+        "! sudo apt-get --assume-yes install libviennacl-dev"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "x7tqqDLHLNDr",
+        "outputId": "dc52fffb-2375-4283-da82-f6327a4d73ad"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "v1.7.0\n"
+          ]
+        }
+      ],
+      "source": [
+        "ver_openpose = \"v1.7.0\"\n",
+        "! echo $ver_openpose"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b11_hkSgLDl5"
+      },
+      "outputs": [],
+      "source": [
+        "! git clone  --depth 1 -b \"$ver_openpose\" https://github.com/CMU-Perceptual-Computing-Lab/openpose.git"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "s-XyxfV8Q-DE"
+      },
+      "outputs": [],
+      "source": [
+        "# manually downloading openpose models\n",
+        "%%bash\n",
+        "gdown 1QCSxJZpnWvM00hx49CJ2zky7PWGzpcEh\n",
+        "unzip models.zip\n",
+        "mv /content/models/face/pose_iter_116000.caffemodel /content/openpose/models/face/pose_iter_116000.caffemodel\n",
+        "mv /content/models/hand/pose_iter_102000.caffemodel /content/openpose/models/hand/pose_iter_102000.caffemodel\n",
+        "mv /content/models/pose/body_25/pose_iter_584000.caffemodel /content/openpose/models/pose/body_25/pose_iter_584000.caffemodel\n",
+        "mv /content/models/pose/coco/pose_iter_440000.caffemodel /content/openpose/models/pose/coco/pose_iter_440000.caffemodel\n",
+        "mv /content/models/pose/mpi/pose_iter_160000.caffemodel /content/openpose/models/pose/mpi/pose_iter_160000.caffemodel\n",
+        "rm -rf models\n",
+        "rm models.zip"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 6,
+      "metadata": {
+        "id": "Bs-zIObzQLYj"
+      },
+      "outputs": [],
+      "source": [
+        "! cd openpose && mkdir build && cd build"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7i7oHh2vQqHv"
+      },
+      "outputs": [],
+      "source": [
+        "! cd openpose/build && cmake -DUSE_CUDNN=OFF -DBUILD_PYTHON=ON .."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "iBvxsDM-EYJk"
+      },
+      "outputs": [],
+      "source": [
+        "# ! cd openpose/build && cmake .."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XEAY8VW0RzD0"
+      },
+      "outputs": [],
+      "source": [
+        "! cd openpose/build && make -j`nproc`\n",
+        "! cd openpose && mkdir output"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 9,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "60nEQBKefg3f",
+        "outputId": "91903854-5dc4-4661-c6d6-cae5ab56bdf2"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Collecting flask-ngrok\n",
+            "  Downloading flask_ngrok-0.0.25-py3-none-any.whl (3.1 kB)\n",
+            "Requirement already satisfied: Flask>=0.8 in /usr/local/lib/python3.10/dist-packages (from flask-ngrok) (2.2.5)\n",
+            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from flask-ngrok) (2.31.0)\n",
+            "Requirement already satisfied: Werkzeug>=2.2.2 in /usr/local/lib/python3.10/dist-packages (from Flask>=0.8->flask-ngrok) (3.0.1)\n",
+            "Requirement already satisfied: Jinja2>=3.0 in /usr/local/lib/python3.10/dist-packages (from Flask>=0.8->flask-ngrok) (3.1.2)\n",
+            "Requirement already satisfied: itsdangerous>=2.0 in /usr/local/lib/python3.10/dist-packages (from Flask>=0.8->flask-ngrok) (2.1.2)\n",
+            "Requirement already satisfied: click>=8.0 in /usr/local/lib/python3.10/dist-packages (from Flask>=0.8->flask-ngrok) (8.1.7)\n",
+            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->flask-ngrok) (3.3.2)\n",
+            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->flask-ngrok) (3.6)\n",
+            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->flask-ngrok) (2.0.7)\n",
+            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->flask-ngrok) (2023.11.17)\n",
+            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from Jinja2>=3.0->Flask>=0.8->flask-ngrok) (2.1.3)\n",
+            "Installing collected packages: flask-ngrok\n",
+            "Successfully installed flask-ngrok-0.0.25\n",
+            "Collecting pyngrok==4.1.1\n",
+            "  Downloading pyngrok-4.1.1.tar.gz (18 kB)\n",
+            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
+            "Requirement already satisfied: future in /usr/local/lib/python3.10/dist-packages (from pyngrok==4.1.1) (0.18.3)\n",
+            "Requirement already satisfied: PyYAML in /usr/local/lib/python3.10/dist-packages (from pyngrok==4.1.1) (6.0.1)\n",
+            "Building wheels for collected packages: pyngrok\n",
+            "  Building wheel for pyngrok (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
+            "  Created wheel for pyngrok: filename=pyngrok-4.1.1-py3-none-any.whl size=15963 sha256=b8ae0d70bccacfc72462262492c96e88942fd5e8113a0a9c6745caea83aad689\n",
+            "  Stored in directory: /root/.cache/pip/wheels/4c/7c/4c/632fba2ea8e88d8890102eb07bc922e1ca8fa14db5902c91a8\n",
+            "Successfully built pyngrok\n",
+            "Installing collected packages: pyngrok\n",
+            "Successfully installed pyngrok-4.1.1\n",
+            "Authtoken saved to configuration file: /root/.ngrok2/ngrok.yml\n"
+          ]
+        }
+      ],
+      "source": [
+        "!pip install flask-ngrok\n",
+        "!pip install pyngrok==4.1.1\n",
+        "!ngrok authtoken <your_token>"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 10,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Fo-q2Q-XMFen",
+        "outputId": "a762d258-7da5-44fc-fdad-7fef43ddd361"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "/content\n",
+            "Cloning into 'clothes-virtual-try-on'...\n",
+            "remote: Enumerating objects: 154, done.\u001b[K\n",
+            "remote: Counting objects: 100% (22/22), done.\u001b[K\n",
+            "remote: Compressing objects: 100% (10/10), done.\u001b[K\n",
+            "remote: Total 154 (delta 16), reused 12 (delta 12), pack-reused 132\u001b[K\n",
+            "Receiving objects: 100% (154/154), 20.47 MiB | 33.87 MiB/s, done.\n",
+            "Resolving deltas: 100% (54/54), done.\n"
+          ]
+        }
+      ],
+      "source": [
+        "import os\n",
+        "%cd /content/\n",
+        "!rm -rf clothes-virtual-try-on\n",
+        "!git clone https://github.com/practice404/clothes-virtual-try-on.git\n",
+        "os.makedirs(\"/content/clothes-virtual-try-on/checkpoints\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 11,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "tnud6ptA9ZwL",
+        "outputId": "bc5ee612-eb57-4118-f5e0-6221484c9571"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "/usr/local/lib/python3.10/dist-packages/gdown/cli.py:126: FutureWarning: Option `--id` was deprecated in version 4.3.1 and will be removed in 5.0. You don't need to pass it anymore to use a file ID.\n",
+            "  warnings.warn(\n",
+            "Downloading...\n",
+            "From (uriginal): https://drive.google.com/uc?id=18q4lS7cNt1_X8ewCgya1fq0dSk93jTL6\n",
+            "From (redirected): https://drive.google.com/uc?id=18q4lS7cNt1_X8ewCgya1fq0dSk93jTL6&confirm=t&uuid=6db7053d-df6d-41e3-ac5b-2f924303335f\n",
+            "To: /content/clothes-virtual-try-on/checkpoints/alias_final.pth\n",
+            "100% 402M/402M [00:01<00:00, 254MB/s]\n",
+            "/usr/local/lib/python3.10/dist-packages/gdown/cli.py:126: FutureWarning: Option `--id` was deprecated in version 4.3.1 and will be removed in 5.0. You don't need to pass it anymore to use a file ID.\n",
+            "  warnings.warn(\n",
+            "Downloading...\n",
+            "From: https://drive.google.com/uc?id=1uDRPY8gh9sHb3UDonq6ZrINqDOd7pmTz\n",
+            "To: /content/clothes-virtual-try-on/checkpoints/gmm_final.pth\n",
+            "100% 76.2M/76.2M [00:00<00:00, 223MB/s]\n",
+            "/usr/local/lib/python3.10/dist-packages/gdown/cli.py:126: FutureWarning: Option `--id` was deprecated in version 4.3.1 and will be removed in 5.0. You don't need to pass it anymore to use a file ID.\n",
+            "  warnings.warn(\n",
+            "Downloading...\n",
+            "From (uriginal): https://drive.google.com/uc?id=1d7lZNLh51Qt5Mi1lXqyi6Asb2ncLrEdC\n",
+            "From (redirected): https://drive.google.com/uc?id=1d7lZNLh51Qt5Mi1lXqyi6Asb2ncLrEdC&confirm=t&uuid=78aeda19-f21d-4598-8bdf-d08a78a99149\n",
+            "To: /content/clothes-virtual-try-on/checkpoints/seg_final.pth\n",
+            "100% 138M/138M [00:01<00:00, 135MB/s]\n"
+          ]
+        }
+      ],
+      "source": [
+        "!gdown --id 18q4lS7cNt1_X8ewCgya1fq0dSk93jTL6 --output /content/clothes-virtual-try-on/checkpoints/alias_final.pth\n",
+        "!gdown --id 1uDRPY8gh9sHb3UDonq6ZrINqDOd7pmTz --output /content/clothes-virtual-try-on/checkpoints/gmm_final.pth\n",
+        "!gdown --id 1d7lZNLh51Qt5Mi1lXqyi6Asb2ncLrEdC --output /content/clothes-virtual-try-on/checkpoints/seg_final.pth"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 12,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "qWPkjShFMK82",
+        "outputId": "cf51a4d3-4833-4788-9878-92a791a944b8"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "/usr/local/lib/python3.10/dist-packages/gdown/cli.py:126: FutureWarning: Option `--id` was deprecated in version 4.3.1 and will be removed in 5.0. You don't need to pass it anymore to use a file ID.\n",
+            "  warnings.warn(\n",
+            "Downloading...\n",
+            "From (uriginal): https://drive.google.com/uc?id=1ysEoAJNxou7RNuT9iKOxRhjVRNY5RLjx\n",
+            "From (redirected): https://drive.google.com/uc?id=1ysEoAJNxou7RNuT9iKOxRhjVRNY5RLjx&confirm=t&uuid=50dc2d49-15b3-47ed-905f-fc2455dfea07\n",
+            "To: /content/clothes-virtual-try-on/cloth_segm_u2net_latest.pth\n",
+            "100% 177M/177M [00:00<00:00, 178MB/s]\n"
+          ]
+        }
+      ],
+      "source": [
+        "!gdown --id 1ysEoAJNxou7RNuT9iKOxRhjVRNY5RLjx --output /content/clothes-virtual-try-on/cloth_segm_u2net_latest.pth --no-cookies"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 13,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "I9MhYntvMP84",
+        "outputId": "774ecb82-0f56-4b65-a6fd-0baa9416d75c"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "/content\n",
+            "Collecting ninja\n",
+            "  Downloading ninja-1.11.1.1-py2.py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl (307 kB)\n",
+            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m307.2/307.2 kB\u001b[0m \u001b[31m2.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+            "\u001b[?25hInstalling collected packages: ninja\n",
+            "Successfully installed ninja-1.11.1.1\n"
+          ]
+        }
+      ],
+      "source": [
+        "%cd /content/\n",
+        "!pip install ninja"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 14,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Rz9LOnvyMWEJ",
+        "outputId": "55af9ee0-cdf5-495f-a14f-ac93695a5fbe"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Cloning into 'Self-Correction-Human-Parsing'...\n",
+            "remote: Enumerating objects: 719, done.\u001b[K\n",
+            "remote: Counting objects: 100% (719/719), done.\u001b[K\n",
+            "remote: Compressing objects: 100% (568/568), done.\u001b[K\n",
+            "remote: Total 719 (delta 149), reused 611 (delta 140), pack-reused 0\u001b[K\n",
+            "Receiving objects: 100% (719/719), 3.88 MiB | 12.81 MiB/s, done.\n",
+            "Resolving deltas: 100% (149/149), done.\n",
+            "/content/Self-Correction-Human-Parsing\n"
+          ]
+        }
+      ],
+      "source": [
+        "!git clone https://github.com/PeikeLi/Self-Correction-Human-Parsing\n",
+        "%cd Self-Correction-Human-Parsing\n",
+        "!mkdir checkpoints"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 15,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "b2k0DLCsMaG0",
+        "outputId": "a28b0d51-14a3-426b-a2cb-d9b209e2b202"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "/usr/local/lib/python3.10/dist-packages/gdown/cli.py:126: FutureWarning: Option `--id` was deprecated in version 4.3.1 and will be removed in 5.0. You don't need to pass it anymore to use a file ID.\n",
+            "  warnings.warn(\n",
+            "Downloading...\n",
+            "From (uriginal): https://drive.google.com/uc?id=1k4dllHpu0bdx38J7H28rVVLpU-kOHmnH\n",
+            "From (redirected): https://drive.google.com/uc?id=1k4dllHpu0bdx38J7H28rVVLpU-kOHmnH&confirm=t&uuid=83091795-9ef5-449e-8d11-008bbe2238eb\n",
+            "To: /content/Self-Correction-Human-Parsing/exp-schp-201908261155-lip.pth\n",
+            "100% 267M/267M [00:01<00:00, 210MB/s]\n"
+          ]
+        }
+      ],
+      "source": [
+        "# downloading LIP dataset model\n",
+        "!gdown --id 1k4dllHpu0bdx38J7H28rVVLpU-kOHmnH\n",
+        "!mv /content/Self-Correction-Human-Parsing/exp-schp-201908261155-lip.pth /content/Self-Correction-Human-Parsing/checkpoints/final.pth"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 16,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "2Y4f3VRyMd9Z",
+        "outputId": "c81b2103-5c6b-4af5-aad1-c51dc65b1015"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "/content\n"
+          ]
+        }
+      ],
+      "source": [
+        "%cd /content/"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "1k2dVj4vMhwA"
+      },
+      "outputs": [],
+      "source": [
+        "%%bash\n",
+        "MINICONDA_INSTALLER_SCRIPT=Miniconda3-4.5.4-Linux-x86_64.sh\n",
+        "MINICONDA_PREFIX=/usr/local\n",
+        "wget https://repo.continuum.io/miniconda/$MINICONDA_INSTALLER_SCRIPT\n",
+        "chmod +x $MINICONDA_INSTALLER_SCRIPT\n",
+        "./$MINICONDA_INSTALLER_SCRIPT -b -f -p $MINICONDA_PREFIX\n",
+        "conda install --channel defaults conda python=3.8 --yes\n",
+        "conda update --channel defaults --all --yes"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 18,
+      "metadata": {
+        "id": "I6entwp3MliV"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "_ = (sys.path\n",
+        "        .append(\"/usr/local/lib/python3.6/site-packages\"))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cseosViyMtYx"
+      },
+      "outputs": [],
+      "source": [
+        "!conda install --channel conda-forge featuretools --yes"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "BEnK6NI6M0cz"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install opencv-python torchgeometry"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "HkySNWttHdW2"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install torchvision"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-wvtaXujRhNp"
+      },
+      "source": [
+        "# Welcome to Virtual-Cloth-Assistant\n",
+        "\n",
+        "> It'll take some extra time in first execution for setting up and downloading of model weights"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 22,
+      "metadata": {
+        "id": "RwcUm39LM8H0"
+      },
+      "outputs": [],
+      "source": [
+        "def make_dir():\n",
+        "  os.system(\"cd /content/ && mkdir inputs && cd inputs && mkdir test && cd test && mkdir cloth cloth-mask image image-parse openpose-img openpose-json\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 28,
+      "metadata": {
+        "id": "9jGRSFuEM9-q"
+      },
+      "outputs": [],
+      "source": [
+        "from flask import Flask, request, send_file, jsonify\n",
+        "from flask_ngrok import run_with_ngrok\n",
+        "from PIL import Image\n",
+        "import base64\n",
+        "import io\n",
+        "\n",
+        "app = Flask(__name__)\n",
+        "run_with_ngrok(app)\n",
+        "\n",
+        "@app.route(\"/\")\n",
+        "def home():\n",
+        "  return jsonify(\"hello world\");\n",
+        "\n",
+        "@app.route(\"/api/transform\", methods=['POST'])\n",
+        "def begin():\n",
+        "  make_dir()\n",
+        "  print(\"data recieved\")\n",
+        "  cloth = request.files['cloth']\n",
+        "  model = request.files['model']\n",
+        "\n",
+        "  cloth = Image.open(cloth.stream)\n",
+        "  model = Image.open(model.stream)\n",
+        "\n",
+        "  cloth.save(\"/content/inputs/test/cloth/cloth.jpg\")\n",
+        "  model.save(\"/content/inputs/test/image/model.jpg\")\n",
+        "\n",
+        "  # running script to compute the predictions\n",
+        "  os.system(\"python /content/clothes-virtual-try-on/run.py\")\n",
+        "\n",
+        "  # loading output\n",
+        "  op = os.listdir(\"/content/output\")[0]\n",
+        "  op = Image.open(f\"/content/output/{op}\")\n",
+        "  buffer = io.BytesIO()\n",
+        "  op.save(buffer, 'png')\n",
+        "  buffer.seek(0)\n",
+        "  os.system(\"rm -rf /content/output/\")\n",
+        "  return send_file(buffer, mimetype='image/gif')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 29,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Pl52gbcqZ3GL",
+        "outputId": "72cbd1c5-afa8-4411-9b3e-75b192b65a07"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            " * Serving Flask app '__main__'\n",
+            " * Debug mode: off\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "INFO:werkzeug:\u001b[31m\u001b[1mWARNING: This is a development server. Do not use it in a production deployment. Use a production WSGI server instead.\u001b[0m\n",
+            " * Running on http://127.0.0.1:5000\n",
+            "INFO:werkzeug:\u001b[33mPress CTRL+C to quit\u001b[0m\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            " * Running on http://e793-34-123-73-186.ngrok-free.app\n",
+            " * Traffic stats available on http://127.0.0.1:4040\n",
+            "data recieved\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "INFO:werkzeug:127.0.0.1 - - [19/Dec/2023 14:36:09] \"POST /api/transform HTTP/1.1\" 200 -\n"
+          ]
+        }
+      ],
+      "source": [
+        "if __name__ == '__main__':\n",
+        "    app.run()"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [
+        "8gqt11Y_RYxU"
+      ],
+      "provenance": []
+    },
+    "gpuClass": "standard",
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

test.py

@@ -0,0 +1,155 @@
+import argparse
+import os
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+import torchgeometry as tgm
+
+from datasets import VITONDataset, VITONDataLoader
+from network import SegGenerator, GMM, ALIASGenerator
+from utils import gen_noise, load_checkpoint, save_images
+
+
+def get_opt():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--name', type=str, required=True)
+
+    parser.add_argument('-b', '--batch_size', type=int, default=1)
+    parser.add_argument('-j', '--workers', type=int, default=1)
+    parser.add_argument('--load_height', type=int, default=1024)
+    parser.add_argument('--load_width', type=int, default=768)
+    parser.add_argument('--shuffle', action='store_true')
+
+    parser.add_argument('--dataset_dir', type=str, default='./datasets/')
+    parser.add_argument('--dataset_mode', type=str, default='test')
+    parser.add_argument('--dataset_list', type=str, default='test_pairs.txt')
+    parser.add_argument('--checkpoint_dir', type=str, default='./checkpoints/')
+    parser.add_argument('--save_dir', type=str, default='./results/')
+
+    parser.add_argument('--display_freq', type=int, default=1)
+
+    parser.add_argument('--seg_checkpoint', type=str, default='seg_final.pth')
+    parser.add_argument('--gmm_checkpoint', type=str, default='gmm_final.pth')
+    parser.add_argument('--alias_checkpoint', type=str, default='alias_final.pth')
+
+    # common
+    parser.add_argument('--semantic_nc', type=int, default=13, help='# of human-parsing map classes')
+    parser.add_argument('--init_type', choices=['normal', 'xavier', 'xavier_uniform', 'kaiming', 'orthogonal', 'none'], default='xavier')
+    parser.add_argument('--init_variance', type=float, default=0.02, help='variance of the initialization distribution')
+
+    # for GMM
+    parser.add_argument('--grid_size', type=int, default=5)
+
+    # for ALIASGenerator
+    parser.add_argument('--norm_G', type=str, default='spectralaliasinstance')
+    parser.add_argument('--ngf', type=int, default=64, help='# of generator filters in the first conv layer')
+    parser.add_argument('--num_upsampling_layers', choices=['normal', 'more', 'most'], default='most',
+                        help='If \'more\', add upsampling layer between the two middle resnet blocks. '
+                             'If \'most\', also add one more (upsampling + resnet) layer at the end of the generator.')
+
+    opt = parser.parse_args()
+    return opt
+
+
+def test(opt, seg, gmm, alias):
+    up = nn.Upsample(size=(opt.load_height, opt.load_width), mode='bilinear')
+    gauss = tgm.image.GaussianBlur((15, 15), (3, 3))
+    gauss.cuda()
+
+    test_dataset = VITONDataset(opt)
+    test_loader = VITONDataLoader(opt, test_dataset)
+
+    with torch.no_grad():
+        for i, inputs in enumerate(test_loader.data_loader):
+            img_names = inputs['img_name']
+            c_names = inputs['c_name']['unpaired']
+
+            img_agnostic = inputs['img_agnostic'].cuda()
+            parse_agnostic = inputs['parse_agnostic'].cuda()
+            pose = inputs['pose'].cuda()
+            c = inputs['cloth']['unpaired'].cuda()
+            cm = inputs['cloth_mask']['unpaired'].cuda()
+
+            # Part 1. Segmentation generation
+            parse_agnostic_down = F.interpolate(parse_agnostic, size=(256, 192), mode='bilinear')
+            pose_down = F.interpolate(pose, size=(256, 192), mode='bilinear')
+            c_masked_down = F.interpolate(c * cm, size=(256, 192), mode='bilinear')
+            cm_down = F.interpolate(cm, size=(256, 192), mode='bilinear')
+            seg_input = torch.cat((cm_down, c_masked_down, parse_agnostic_down, pose_down, gen_noise(cm_down.size()).cuda()), dim=1)
+
+            parse_pred_down = seg(seg_input)
+            parse_pred = gauss(up(parse_pred_down))
+            parse_pred = parse_pred.argmax(dim=1)[:, None]
+
+            parse_old = torch.zeros(parse_pred.size(0), 13, opt.load_height, opt.load_width, dtype=torch.float).cuda()
+            parse_old.scatter_(1, parse_pred, 1.0)
+
+            labels = {
+                0:  ['background',  [0]],
+                1:  ['paste',       [2, 4, 7, 8, 9, 10, 11]],
+                2:  ['upper',       [3]],
+                3:  ['hair',        [1]],
+                4:  ['left_arm',    [5]],
+                5:  ['right_arm',   [6]],
+                6:  ['noise',       [12]]
+            }
+            parse = torch.zeros(parse_pred.size(0), 7, opt.load_height, opt.load_width, dtype=torch.float).cuda()
+            for j in range(len(labels)):
+                for label in labels[j][1]:
+                    parse[:, j] += parse_old[:, label]
+
+            # Part 2. Clothes Deformation
+            agnostic_gmm = F.interpolate(img_agnostic, size=(256, 192), mode='nearest')
+            parse_cloth_gmm = F.interpolate(parse[:, 2:3], size=(256, 192), mode='nearest')
+            pose_gmm = F.interpolate(pose, size=(256, 192), mode='nearest')
+            c_gmm = F.interpolate(c, size=(256, 192), mode='nearest')
+            gmm_input = torch.cat((parse_cloth_gmm, pose_gmm, agnostic_gmm), dim=1)
+
+            _, warped_grid = gmm(gmm_input, c_gmm)
+            warped_c = F.grid_sample(c, warped_grid, padding_mode='border')
+            warped_cm = F.grid_sample(cm, warped_grid, padding_mode='border')
+
+            # Part 3. Try-on synthesis
+            misalign_mask = parse[:, 2:3] - warped_cm
+            misalign_mask[misalign_mask < 0.0] = 0.0
+            parse_div = torch.cat((parse, misalign_mask), dim=1)
+            parse_div[:, 2:3] -= misalign_mask
+
+            output = alias(torch.cat((img_agnostic, pose, warped_c), dim=1), parse, parse_div, misalign_mask)
+
+            unpaired_names = []
+            for img_name, c_name in zip(img_names, c_names):
+                unpaired_names.append('{}_{}'.format(img_name.split('_')[0], c_name))
+
+            save_images(output, unpaired_names, os.path.join(opt.save_dir, opt.name))
+
+            if (i + 1) % opt.display_freq == 0:
+                print("step: {}".format(i + 1))
+
+
+def main():
+    opt = get_opt()
+    print(opt)
+
+    if not os.path.exists(os.path.join(opt.save_dir, opt.name)):
+        os.makedirs(os.path.join(opt.save_dir, opt.name))
+
+    seg = SegGenerator(opt, input_nc=opt.semantic_nc + 8, output_nc=opt.semantic_nc)
+    gmm = GMM(opt, inputA_nc=7, inputB_nc=3)
+    opt.semantic_nc = 7
+    alias = ALIASGenerator(opt, input_nc=9)
+    opt.semantic_nc = 13
+
+    load_checkpoint(seg, os.path.join(opt.checkpoint_dir, opt.seg_checkpoint))
+    load_checkpoint(gmm, os.path.join(opt.checkpoint_dir, opt.gmm_checkpoint))
+    load_checkpoint(alias, os.path.join(opt.checkpoint_dir, opt.alias_checkpoint))
+
+    seg.cuda().eval()
+    gmm.cuda().eval()
+    alias.cuda().eval()
+    test(opt, seg, gmm, alias)
+
+
+if __name__ == '__main__':
+    main()

utils.py

@@ -0,0 +1,40 @@
+import os
+
+import cv2
+import numpy as np
+from PIL import Image
+import torch
+
+
+def gen_noise(shape):
+    noise = np.zeros(shape, dtype=np.uint8)
+    ### noise
+    noise = cv2.randn(noise, 0, 255)
+    noise = np.asarray(noise / 255, dtype=np.uint8)
+    noise = torch.tensor(noise, dtype=torch.float32)
+    return noise
+
+
+def save_images(img_tensors, img_names, save_dir):
+    for img_tensor, img_name in zip(img_tensors, img_names):
+        tensor = (img_tensor.clone()+1)*0.5 * 255
+        tensor = tensor.cpu().clamp(0,255)
+
+        try:
+            array = tensor.numpy().astype('uint8')
+        except:
+            array = tensor.detach().numpy().astype('uint8')
+
+        if array.shape[0] == 1:
+            array = array.squeeze(0)
+        elif array.shape[0] == 3:
+            array = array.swapaxes(0, 1).swapaxes(1, 2)
+
+        im = Image.fromarray(array)
+        im.save(os.path.join(save_dir, img_name), format='JPEG')
+
+
+def load_checkpoint(model, checkpoint_path):
+    if not os.path.exists(checkpoint_path):
+        raise ValueError("'{}' is not a valid checkpoint path".format(checkpoint_path))
+    model.load_state_dict(torch.load(checkpoint_path))