init_commit

.gitignore

@@ -0,0 +1,130 @@
+# 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/
+pip-wheel-metadata/
+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/
+
+# 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
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.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
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__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/
+.vercel

Dockerfile

@@ -0,0 +1,32 @@
+# read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
+# you will also find guides on how best to write your Dockerfile
+
+# Include base image
+FROM docker.io/pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
+
+# Define working directory
+WORKDIR /workspace/
+
+# Set timezone
+ENV TZ=Asia/Tokyo
+RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
+
+# Install dependencies
+RUN apt-get update && apt-get -y install libgl1 libglib2.0-0 vim
+RUN apt-get autoremove -y && apt-get clean -y
+
+# Add pretrained model
+ADD seg2art ./seg2art
+ADD static ./static
+ADD templates ./templates
+ADD utils ./utils
+
+# Add necessary files
+ADD app.py ./
+
+# pip install
+ADD requirements.txt ./
+RUN pip install -r requirements.txt
+
+# Run server
+CMD [ "python", "-u", "./app.py" ]

app.py

@@ -0,0 +1,130 @@
+import os
+import sys
+import time
+import json
+import torch
+import base64
+from PIL import Image
+from io import BytesIO
+
+# set CUDA_MODULE_LOADING=LAZY to speed up the serverless function
+os.environ["CUDA_MODULE_LOADING"] = "LAZY"
+# set SAFETENSORS_FAST_GPU=1 to speed up the serverless function
+os.environ["SAFETENSORS_FAST_GPU"] = "1"
+
+sys.path.append(os.path.join(os.path.dirname(__file__), "seg2art"))
+from seg2art.sstan_models.pix2pix_model import Pix2PixModel
+from seg2art.options.test_options import TestOptions
+from seg2art.inference_util import get_artwork
+
+import uvicorn
+from fastapi import FastAPI, Form
+from fastapi.templating import Jinja2Templates
+from fastapi.responses import PlainTextResponse, HTMLResponse
+from fastapi.requests import Request
+from fastapi.staticfiles import StaticFiles
+
+
+# declare constants
+HOST = "0.0.0.0"
+PORT = 7860
+# FastAPI
+app = FastAPI(root_path=os.path.abspath(os.path.dirname(__file__)))
+app.mount("/static", StaticFiles(directory="static"), name="static")
+templates = Jinja2Templates(directory="templates")
+
+
+# initialize SEAN model.
+opt = TestOptions().parse()
+opt.status = "test"
+model = Pix2PixModel(opt)
+model = model.half() if torch.cuda.is_available() else model
+model.eval()
+
+
+from utils.umap_utils import get_code, load_boundries, modify_code
+
+boundaries = load_boundries()
+global current_codes
+current_codes = {}
+max_user_num = 5
+
+initial_code_path = os.path.join(os.path.dirname(__file__), "static/init_code")
+initial_code = torch.load(initial_code_path) if torch.cuda.is_available() else torch.load(initial_code_path, map_location=torch.device("cpu"))
+
+
+def EncodeImage(img_pil):
+    with BytesIO() as buffer:
+        img_pil.save(buffer, "jpeg")
+        image_data = base64.b64encode(buffer.getvalue())
+    return image_data
+
+
+def DecodeImage(img_pil):
+    img_pil = BytesIO(base64.urlsafe_b64decode(img_pil))
+    img_pil = Image.open(img_pil).convert("RGB")
+    return img_pil
+
+
+def process_input(body, random=False):
+    global current_codes
+    json_body = json.loads(body.decode("utf-8"))
+    user_id = json_body["user_id"]
+    start_time = time.time()
+
+    # save current code for different users
+    if user_id not in current_codes:
+        current_codes[user_id] = initial_code.clone()
+    if len(current_codes[user_id]) > max_user_num:
+        current_codes[user_id] = current_codes[user_id][-max_user_num:]
+
+    if random:
+        # randomize code
+        domain = json_body["model"]
+        current_codes[user_id] = get_code(domain, boundaries)
+
+    # get input
+    input_img = DecodeImage(json_body["img"])
+
+    try:
+        move_range = float(json_body["move_range"])
+    except:
+        move_range = 0
+
+    # set move range to 3 if random is True
+    move_range = 3 if random else move_range
+    # print("Input image was received")
+    # get selected style
+    domain = json_body["model"]
+    if move_range != 0:
+        modified_code = modify_code(current_codes[user_id], boundaries, domain, move_range)
+    else:
+        modified_code = current_code.clone()
+
+    # inference
+    result = get_artwork(model, input_img, modified_code)
+    print("Time Cost: ", time.time() - start_time)
+    return EncodeImage(result)
+
+
+@app.get("/", response_class=HTMLResponse)
+def root(request: Request):
+    return templates.TemplateResponse("index.html", {"request": request})
+
+
+@app.post("/predict")
+async def predict(request: Request):
+    body = await request.body()
+    result = process_input(body, random=False)
+    return result
+
+
+@app.post("/predict_random")
+async def predict_random(request: Request):
+    body = await request.body()
+    result = process_input(body, random=True)
+    return result
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host=HOST, port=PORT, log_level="info")

requirements.txt

@@ -0,0 +1,6 @@
+scikit-learn
+scikit-image
+torchvision>=0.7.0
+torch>=1.6.0
+fastapi
+uvicorn

seg2art/checkpoints/multimodal_artworks/latest_net_G-fp16.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:11e5d324b59dce20e81cea0eed77b25c7b9f6b56ccb44970b67593b4287ddb4c
+size 418576205

seg2art/inference_util.py

@@ -0,0 +1,77 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+import numpy as np
+from PIL import Image
+from torchvision import transforms
+
+# define constants
+image_size = 256
+
+# to label
+values = [12, 2, 6, 8, 1, 10, 3, 14, 11, 4, 5, 13, 9]
+values = np.array(values)
+
+# from color
+colors = [
+    (135, 206, 235),
+    (155, 118, 83),
+    (176, 212, 155),
+    (90, 188, 216),
+    (193, 190, 186),
+    (90, 77, 65),
+    (86, 125, 70),
+    (66, 105, 47),
+    (21, 119, 190),
+    (58, 46, 39),
+    (77, 65, 90),
+    (253, 218, 22),
+    (208, 204, 204),
+]
+colors = np.array(colors)
+
+
+def remap_label(arr):
+    # compare only first 1 channel to speed up
+    arr_r = arr[:, :, 0]
+
+    # remap color to label
+    for i in range(len(colors)):
+        arr_r[arr_r == colors[i][0]] = values[i]
+    # others to 15
+    arr_r[arr_r > 15] = 15
+    return arr_r
+
+
+preprocess = transforms.Compose(
+    [
+        transforms.Resize([image_size, image_size]),
+        transforms.ToTensor(),
+    ]
+)
+
+
+def image_loader(loader, label_inp):
+    image = Image.fromarray(label_inp).convert("RGB")
+    image = image.resize((image_size, image_size))
+    image = loader(image).float() * 255
+    image = image.clone().detach().requires_grad_(True)
+    image = image.unsqueeze(0)
+    return image
+
+
+def tensor2im(image_tensor):
+    image_numpy = image_tensor[0].detach().cpu().float().numpy()
+    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
+    image_numpy = np.clip(image_numpy, 0, 255)
+    return Image.fromarray(image_numpy.astype(np.uint8))
+
+
+def get_artwork(model, data, code):
+    label_inp = remap_label(np.array(data))
+    label_inp = (image_loader(preprocess, label_inp)).detach().half()
+
+    image_out = model(label_inp, mode="inference", style_codes=code)
+    image_out = tensor2im(image_out)
+    return image_out

seg2art/model_util.py

@@ -0,0 +1,158 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import re
+import importlib
+import torch
+from argparse import Namespace
+import numpy as np
+from PIL import Image
+import os
+
+
+# Converts a Tensor into a Numpy array
+# |imtype|: the desired type of the converted numpy array
+def tensor2im(image_tensor, imtype=np.uint8, normalize=True, tile=False):
+    if isinstance(image_tensor, list):
+        image_numpy = []
+        for i in range(len(image_tensor)):
+            image_numpy.append(tensor2im(image_tensor[i], imtype, normalize))
+        return image_numpy
+
+    if image_tensor.dim() == 4:
+        # transform each image in the batch
+        images_np = []
+        for b in range(image_tensor.size(0)):
+            one_image = image_tensor[b]
+            one_image_np = tensor2im(one_image)
+            images_np.append(one_image_np.reshape(1, *one_image_np.shape))
+        images_np = np.concatenate(images_np, axis=0)
+        return images_np
+
+    if image_tensor.dim() == 2:
+        image_tensor = image_tensor.unsqueeze(0)
+    image_numpy = image_tensor.detach().cpu().float().numpy()
+    if normalize:
+        image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
+    else:
+        image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0
+    image_numpy = np.clip(image_numpy, 0, 255)
+    if image_numpy.shape[2] == 1:
+        image_numpy = image_numpy[:, :, 0]
+    return image_numpy.astype(imtype)
+
+
+# Converts a one-hot tensor into a colorful label map
+def tensor2label(label_tensor, n_label, imtype=np.uint8, tile=False):
+    if label_tensor.dim() == 4:
+        # transform each image in the batch
+        images_np = []
+        for b in range(label_tensor.size(0)):
+            one_image = label_tensor[b]
+            one_image_np = tensor2label(one_image, n_label, imtype)
+            images_np.append(one_image_np.reshape(1, *one_image_np.shape))
+        images_np = np.concatenate(images_np, axis=0)
+        if tile:
+            images_tiled = tile_images(images_np)
+            return images_tiled
+        else:
+            images_np = images_np[0]
+            return images_np
+
+    if label_tensor.dim() == 1:
+        return np.zeros((64, 64, 3), dtype=np.uint8)
+    if n_label == 0:
+        return tensor2im(label_tensor, imtype)
+    label_tensor = label_tensor.cpu().float()
+    if label_tensor.size()[0] > 1:
+        label_tensor = label_tensor.max(0, keepdim=True)[1]
+    label_tensor = Colorize(n_label)(label_tensor)
+    label_numpy = np.transpose(label_tensor.numpy(), (1, 2, 0))
+    result = label_numpy.astype(imtype)
+    return result
+
+
+def save_image(image_numpy, image_path, create_dir=False):
+    if create_dir:
+        os.makedirs(os.path.dirname(image_path), exist_ok=True)
+    if len(image_numpy.shape) == 2:
+        image_numpy = np.expand_dims(image_numpy, axis=2)
+    if image_numpy.shape[2] == 1:
+        image_numpy = np.repeat(image_numpy, 3, 2)
+    image_pil = Image.fromarray(image_numpy)
+
+    # save to png
+    image_pil.save(image_path.replace('.jpg', '.png'))
+
+
+def mkdirs(paths):
+    if isinstance(paths, list) and not isinstance(paths, str):
+        for path in paths:
+            mkdir(path)
+    else:
+        mkdir(paths)
+
+
+def mkdir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split('(\d+)', text)]
+
+
+def natural_sort(items):
+    items.sort(key=natural_keys)
+
+
+def str2bool(v):
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+
+def find_class_in_module(target_cls_name, module):
+    target_cls_name = target_cls_name.replace('_', '').lower()
+    clslib = importlib.import_module(module)
+    cls = None
+    for name, clsobj in clslib.__dict__.items():
+        if name.lower() == target_cls_name:
+            cls = clsobj
+
+    if cls is None:
+        print("In %s, there should be a class whose name matches %s in lowercase without underscore(_)" % (module, target_cls_name))
+        exit(0)
+
+    return cls
+
+
+def save_network(net, label, epoch, opt):
+    save_filename = '%s_net_%s.pth' % (epoch, label)
+    save_path = os.path.join(opt.checkpoints_dir, opt.name, save_filename)
+    torch.save(net.cpu().state_dict(), save_path)
+    if len(opt.gpu_ids) and torch.cuda.is_available():
+        net.cuda()
+
+
+def load_network(net, label, epoch, opt):
+    save_filename = '%s_net_%s.pth' % (epoch, label)
+    save_dir = os.path.join(opt.checkpoints_dir, opt.name)
+    save_path = os.path.join(save_dir, save_filename)
+    weights = torch.load(save_path)
+    net.load_state_dict(weights, strict=False)#
+    return net

seg2art/options/__init__.py

@@ -0,0 +1,4 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""

seg2art/options/base_options.py

@@ -0,0 +1,184 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import os
+import sys
+import torch
+import pickle
+import argparse
+import sstan_models
+import utils as util
+
+
+class BaseOptions():
+    def __init__(self):
+        self.initialized = False
+
+    def initialize(self, parser):
+        # experiment specifics
+        parser.add_argument('--name', type=str, default='multimodal_artworks', help='name of the experiment. It decides where to store samples and sstan_models')
+
+        parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')
+        parser.add_argument('--checkpoints_dir', type=str, default='./seg2art/checkpoints', help='sstan_models are saved here')
+        parser.add_argument('--model', type=str, default='pix2pix', help='which model to use')
+        parser.add_argument('--norm_G', type=str, default='spectralinstance', help='instance normalization or batch normalization')
+        parser.add_argument('--norm_D', type=str, default='spectralinstance', help='instance normalization or batch normalization')
+        parser.add_argument('--norm_E', type=str, default='spectralinstance', help='instance normalization or batch normalization')
+        parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')
+
+        # input/output sizes
+        parser.add_argument('--batchSize', type=int, default=1, help='input batch size')
+        parser.add_argument('--preprocess_mode', type=str, default='scale_width_and_crop', help='scaling and cropping of images at load time.', choices=("resize_and_crop", "crop", "scale_width", "scale_width_and_crop", "scale_shortside", "scale_shortside_and_crop", "fixed", "none"))
+        parser.add_argument('--load_size', type=int, default=512, help='Scale images to this size. The final image will be cropped to --crop_size.')
+        parser.add_argument('--crop_size', type=int, default=512, help='Crop to the width of crop_size (after initially scaling the images to load_size.)')
+        parser.add_argument('--aspect_ratio', type=float, default=1.0, help='The ratio width/height. The final height of the load image will be crop_size/aspect_ratio')
+        parser.add_argument('--label_nc', type=int, default=16, help='# of input label classes without unknown class. If you have unknown class as class label, specify --contain_dopntcare_label.')
+        parser.add_argument('--contain_dontcare_label', action='store_true', help='if the label map contains dontcare label (dontcare=255)')
+        parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels')
+
+        # for setting inputs
+        parser.add_argument('--dataroot', type=str, default='./datasets/cityscapes/')
+        parser.add_argument('--dataset_mode', type=str, default='custom')
+        parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')
+        parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data argumentation')
+        parser.add_argument('--nThreads', default=0, type=int, help='# threads for loading data')
+        parser.add_argument('--max_dataset_size', type=int, default=sys.maxsize, help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')
+        parser.add_argument('--load_from_opt_file', action='store_true', help='load the options from checkpoints and use that as default')
+        parser.add_argument('--cache_filelist_write', action='store_true', help='saves the current filelist into a text file, so that it loads faster')
+        parser.add_argument('--cache_filelist_read', action='store_true', help='reads from the file list cache')
+
+        # for displays
+        parser.add_argument('--display_winsize', type=int, default=400, help='display window size')
+
+        # for generator
+        parser.add_argument('--netG', type=str, default='spade', help='selects model to use for netG (pix2pixhd | spade)')
+        parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer')
+        parser.add_argument('--init_type', type=str, default='xavier', help='network initialization [normal|xavier|kaiming|orthogonal]')
+        parser.add_argument('--init_variance', type=float, default=0.02, help='variance of the initialization distribution')
+        parser.add_argument('--z_dim', type=int, default=256,
+                            help="dimension of the latent z vector")
+
+        # for instance-wise features
+        parser.add_argument('--no_instance', action='store_true', help='if specified, do *not* add instance map as input')
+        parser.add_argument('--nef', type=int, default=16, help='# of encoder filters in the first conv layer')
+        parser.add_argument('--use_vae', action='store_true', help='enable training with an image encoder.')
+
+        self.initialized = True
+        return parser
+
+    def gather_options(self):
+        # initialize parser with basic options
+        if not self.initialized:
+            parser = argparse.ArgumentParser(
+                formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+            parser = self.initialize(parser)
+
+        # get the basic options
+        opt, unknown = parser.parse_known_args()
+
+        # modify model-related parser options
+        model_name = opt.model
+        model_option_setter = sstan_models.get_option_setter(model_name)
+        parser = model_option_setter(parser, self.isTrain)
+
+
+        # # modify dataset-related parser options
+        # dataset_mode = opt.dataset_mode
+        # dataset_option_setter = data.get_option_setter(dataset_mode)
+        # parser = dataset_option_setter(parser, self.isTrain)
+
+        # opt, unknown = parser.parse_known_args()
+
+        # # if there is opt_file, load it.
+        # # The previous default options will be overwritten
+        # if opt.load_from_opt_file:
+        #     parser = self.update_options_from_file(parser, opt)
+
+        opt = parser.parse_args()
+
+        opt.contain_dontcare_label = False
+        opt.no_instance = True
+        opt.use_vae = False
+
+        self.parser = parser
+        return opt
+
+    def print_options(self, opt):
+        message = ''
+        message += '----------------- Options ---------------\n'
+        for k, v in sorted(vars(opt).items()):
+            comment = ''
+            default = self.parser.get_default(k)
+            if v != default:
+                comment = '\t[default: %s]' % str(default)
+            message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
+        message += '----------------- End -------------------'
+        print(message)
+
+    def option_file_path(self, opt, makedir=False):
+        expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
+        if makedir:
+            util.mkdirs(expr_dir)
+        file_name = os.path.join(expr_dir, 'opt')
+        return file_name
+
+    def save_options(self, opt):
+        file_name = self.option_file_path(opt, makedir=True)
+        with open(file_name + '.txt', 'wt') as opt_file:
+            for k, v in sorted(vars(opt).items()):
+                comment = ''
+                default = self.parser.get_default(k)
+                if v != default:
+                    comment = '\t[default: %s]' % str(default)
+                opt_file.write('{:>25}: {:<30}{}\n'.format(str(k), str(v), comment))
+
+        with open(file_name + '.pkl', 'wb') as opt_file:
+            pickle.dump(opt, opt_file)
+
+    def update_options_from_file(self, parser, opt):
+        new_opt = self.load_options(opt)
+        for k, v in sorted(vars(opt).items()):
+            if hasattr(new_opt, k) and v != getattr(new_opt, k):
+                new_val = getattr(new_opt, k)
+                parser.set_defaults(**{k: new_val})
+        return parser
+
+    def load_options(self, opt):
+        file_name = self.option_file_path(opt, makedir=False)
+        new_opt = pickle.load(open(file_name + '.pkl', 'rb'))
+        return new_opt
+
+    def parse(self, save=False):
+
+        opt = self.gather_options()
+        opt.isTrain = self.isTrain   # train or test
+
+        #self.print_options(opt)
+        if opt.isTrain:
+            self.save_options(opt)
+
+        # Set semantic_nc based on the option.
+        # This will be convenient in many places
+        opt.semantic_nc = opt.label_nc + \
+            (1 if opt.contain_dontcare_label else 0) + \
+            (0 if opt.no_instance else 1)
+
+        # set gpu ids
+        str_ids = opt.gpu_ids.split(',')
+        opt.gpu_ids = []
+        for str_id in str_ids:
+            id = int(str_id)
+            if id >= 0:
+                opt.gpu_ids.append(id)
+        opt.gpu_ids = [] if torch.cuda.device_count() == 0 else opt.gpu_ids
+        if len(opt.gpu_ids) > 0:
+            torch.cuda.set_device(opt.gpu_ids[0])
+
+        assert len(opt.gpu_ids) == 0 or opt.batchSize % len(opt.gpu_ids) == 0, \
+            "Batch size %d is wrong. It must be a multiple of # GPUs %d." \
+            % (opt.batchSize, len(opt.gpu_ids))
+
+        self.opt = opt
+        return self.opt

seg2art/options/test_options.py

@@ -0,0 +1,22 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+from .base_options import BaseOptions
+
+
+class TestOptions(BaseOptions):
+    def initialize(self, parser):
+        BaseOptions.initialize(self, parser)
+        parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.')
+        parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
+        parser.add_argument('--checkpoint_path', type=str, default='./checkpoints/multimodal_artworks/latest_net_G-fp16.pth', help='load model from a checkpoint')
+        parser.add_argument('--how_many', type=int, default=float("inf"), help='how many test images to run')
+
+        parser.set_defaults(preprocess_mode='scale_width_and_crop', crop_size=512, load_size=512, display_winsize=256)
+        parser.set_defaults(serial_batches=True)
+        parser.set_defaults(no_flip=True)
+        parser.set_defaults(phase='test')
+        self.isTrain = False
+        return parser

seg2art/sstan_models/__init__.py

@@ -0,0 +1,44 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import importlib
+import torch
+
+
+def find_model_using_name(model_name):
+    # Given the option --model [modelname],
+    # the file "sstan_models/modelname_model.py"
+    # will be imported.
+    model_filename = "sstan_models." + model_name + "_model"
+    modellib = importlib.import_module(model_filename)
+
+    # In the file, the class called ModelNameModel() will
+    # be instantiated. It has to be a subclass of torch.nn.Module,
+    # and it is case-insensitive.
+    model = None
+    target_model_name = model_name.replace('_', '') + 'model'
+    for name, cls in modellib.__dict__.items():
+        if name.lower() == target_model_name.lower() \
+           and issubclass(cls, torch.nn.Module):
+            model = cls
+
+    if model is None:
+        print("In %s.py, there should be a subclass of torch.nn.Module with class name that matches %s in lowercase." % (model_filename, target_model_name))
+        exit(0)
+
+    return model
+
+
+def get_option_setter(model_name):
+    model_class = find_model_using_name(model_name)
+    return model_class.modify_commandline_options
+
+
+def create_model(opt):
+    model = find_model_using_name(opt.model)
+    instance = model(opt)
+    print("model [%s] was created" % (type(instance).__name__))
+
+    return instance

seg2art/sstan_models/networks/__init__.py

@@ -0,0 +1,63 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import torch
+from sstan_models.networks.base_network import BaseNetwork
+# from sstan_models.networks.loss import *
+# from sstan_models.networks.discriminator import *
+from sstan_models.networks.generator import *
+# from sstan_models.networks.encoder import *
+import model_util as util
+
+
+def find_network_using_name(target_network_name, filename):
+    target_class_name = target_network_name + filename
+    module_name = 'sstan_models.networks.' + filename
+    network = util.find_class_in_module(target_class_name, module_name)
+
+    assert issubclass(network, BaseNetwork), \
+        "Class %s should be a subclass of BaseNetwork" % network
+
+    return network
+
+
+def modify_commandline_options(parser, is_train):
+    opt, _ = parser.parse_known_args()
+
+    netG_cls = find_network_using_name(opt.netG, 'generator')
+    parser = netG_cls.modify_commandline_options(parser, is_train)
+    if is_train:
+        netD_cls = find_network_using_name(opt.netD, 'discriminator')
+        parser = netD_cls.modify_commandline_options(parser, is_train)
+    # netE_cls = find_network_using_name('conv', 'encoder')
+    # parser = netE_cls.modify_commandline_options(parser, is_train)
+
+    return parser
+
+
+def create_network(cls, opt):
+    net = cls(opt)
+    net.print_network()
+    if len(opt.gpu_ids) > 0:
+        assert(torch.cuda.is_available())
+        net.cuda()
+    net.init_weights(opt.init_type, opt.init_variance)
+    return net
+
+
+def define_G(opt):
+    netG_cls = find_network_using_name(opt.netG, 'generator')
+    return create_network(netG_cls, opt)
+
+
+def define_D(opt):
+    netD_cls = find_network_using_name(opt.netD, 'discriminator')
+    return create_network(netD_cls, opt)
+
+
+def define_E(opt):
+    # there exists only one encoder type
+    netE_cls = find_network_using_name('conv', 'encoder')
+    return create_network(netE_cls, opt)

seg2art/sstan_models/networks/architecture.py

@@ -0,0 +1,231 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+import torch.nn.utils.spectral_norm as spectral_norm
+from sstan_models.networks.normalization import SPADE
+
+
+# ResNet block that uses SPADE.
+# It differs from the ResNet block of pix2pixHD in that
+# it takes in the segmentation map as input, learns the skip connection if necessary,
+# and applies normalization first and then convolution.
+# This architecture seemed like a standard architecture for unconditional or
+# class-conditional GAN architecture using residual block.
+# The code was inspired from https://github.com/LMescheder/GAN_stability.
+class SPADEResnetBlock(nn.Module):
+    def __init__(self, fin, fout, opt, feed_code=False):
+        super().__init__()
+
+        self.status = 'train'
+        # Attributes
+        self.learned_shortcut = (fin != fout)
+        fmiddle = min(fin, fout)
+
+        # create conv layers
+        self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=1)
+        self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=1)
+        if self.learned_shortcut:
+            self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False)
+
+        # apply spectral norm if specified
+        if 'spectral' in opt.norm_G:
+            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)
+
+        # define normalization layers
+        spade_config_str = opt.norm_G.replace('spectral', '')
+
+        # Attention related
+        self.channelAtt = ChannelAttention(fout)
+        self.spatialAtt = SpatialAttention()
+
+        self.norm_0 = SPADE(spade_config_str, fin, feed_code, status=self.status, spade_params=[
+                            spade_config_str, fin, opt.semantic_nc])
+        self.norm_1 = SPADE(spade_config_str, fmiddle, feed_code, status=self.status, spade_params=[
+                            spade_config_str, fmiddle, opt.semantic_nc])
+        if self.learned_shortcut:
+            self.norm_s = SPADE(spade_config_str, fin, feed_code, status=self.status, spade_params=[
+                                spade_config_str, fin, opt.semantic_nc])
+
+    # note the resnet block with SPADE also takes in |seg|,
+    # the semantic segmentation map as input
+    def forward(self, x, seg, style_codes=None, self_attention=False):
+        x_s = self.shortcut(x, seg, style_codes=style_codes)
+
+        dx = self.conv_0(self.actvn(
+            self.norm_0(x, seg, style_codes=style_codes)))
+        dx = self.conv_1(self.actvn(self.norm_1(
+            dx, seg, style_codes=style_codes)))
+
+        dx = self.channelAtt(dx) * dx
+        dx = self.spatialAtt(dx) * dx
+        out = x_s + dx
+
+        return out
+
+    def shortcut(self, x, seg, style_codes):
+        if self.learned_shortcut:
+            x_s = self.conv_s(self.norm_s(x, seg, style_codes=style_codes))
+        else:
+            x_s = x
+        return x_s
+
+    def actvn(self, x):
+        return F.leaky_relu(x, 2e-1)
+
+
+# ResNet block used in pix2pixHD
+# We keep the same architecture as pix2pixHD.
+class ResnetBlock(nn.Module):
+    def __init__(self, dim, norm_layer, activation=nn.ReLU(False), kernel_size=3):
+        super().__init__()
+
+        pw = (kernel_size - 1) // 2
+        self.conv_block = nn.Sequential(
+            nn.ReflectionPad2d(pw),
+            norm_layer(nn.Conv2d(dim, dim, kernel_size=kernel_size)),
+            activation,
+            nn.ReflectionPad2d(pw),
+            norm_layer(nn.Conv2d(dim, dim, kernel_size=kernel_size))
+        )
+
+    def forward(self, x):
+        y = self.conv_block(x)
+        out = x + y
+        return out
+
+
+# VGG architecter, used for the perceptual loss using a pretrained VGG network
+class VGG19(torch.nn.Module):
+    def __init__(self, requires_grad=False):
+        super().__init__()
+        vgg_pretrained_features = torchvision.models.vgg19(
+            pretrained=True).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        for x in range(2):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(2, 7):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(7, 12):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(12, 21):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(21, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+        # torch.cuda.empty_cache()
+    def forward(self, X):
+        with torch.cuda.amp.autocast():
+            # with torch.no_grad():
+            h_relu1 = self.slice1(X)
+            h_relu2 = self.slice2(h_relu1)
+            h_relu3 = self.slice3(h_relu2)
+            h_relu4 = self.slice4(h_relu3)
+            h_relu5 = self.slice5(h_relu4)
+            out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
+        return out
+
+'''
+class SourceReferenceAttention(nn.Module):
+    """
+    Source-Reference Attention Layer
+    """
+
+    def __init__(self, in_planes_s, in_planes_r):
+        """
+        Parameters
+        ----------
+            in_planes_s: int
+                Number of input source feature vector channels.
+            in_planes_r: int
+                Number of input reference feature vector channels.
+        """
+        super(SourceReferenceAttention, self).__init__()
+        self.query_conv = nn.Conv2d(in_channels=in_planes_s,
+                                    out_channels=in_planes_s//8, kernel_size=1)
+        self.key_conv = nn.Conv2d(in_channels=in_planes_r,
+                                  out_channels=in_planes_r//8, kernel_size=1)
+        self.value_conv = nn.Conv2d(in_channels=in_planes_r,
+                                    out_channels=in_planes_r, kernel_size=1)
+        self.gamma = nn.Parameter(torch.zeros(1))
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, source, reference):
+        """
+        Parameters
+        ----------
+            source : torch.Tensor
+                Source feature maps (B x Cs x Ts x Hs x Ws)
+            reference : torch.Tensor
+                Reference feature maps (B x Cr x Tr x Hr x Wr )
+         Returns :
+            torch.Tensor
+                Source-reference attention value added to the input source features
+            torch.Tensor
+                Attention map (B x Ns x Nt) (Ns=Ts*Hs*Ws, Nr=Tr*Hr*Wr)
+        """
+        s_batchsize, sC, sH, sW = source.size()
+        r_batchsize, rC, rH, rW = reference.size()
+        proj_query = self.query_conv(source).view(
+            s_batchsize, -1, sH*sW).permute(0, 2, 1)
+        proj_key = self.key_conv(reference).view(r_batchsize, -1, rW*rH)
+        energy = torch.bmm(proj_query, proj_key)
+        attention = self.softmax(energy)
+        proj_value = self.value_conv(reference).view(r_batchsize, -1, rH*rW)
+        out = torch.bmm(proj_value, attention.permute(0, 2, 1))
+        out = out.view(s_batchsize, sC, sH, sW)
+        out = self.gamma*out + source
+        return out, attention
+'''
+
+class ChannelAttention(nn.Module):
+    def __init__(self, in_planes, ratio=16):
+        super(ChannelAttention, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.max_pool = nn.AdaptiveMaxPool2d(1)
+
+        self.fc1 = nn.Conv2d(in_planes, in_planes // 16, 1, bias=False)
+        self.relu1 = nn.ReLU()
+        self.fc2 = nn.Conv2d(in_planes // 16, in_planes, 1, bias=False)
+
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+
+        avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
+        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
+        out = avg_out + max_out
+        return self.sigmoid(out)
+
+
+class SpatialAttention(nn.Module):
+    def __init__(self, kernel_size=7):
+        super(SpatialAttention, self).__init__()
+
+        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
+        padding = 3 if kernel_size == 7 else 1
+
+        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+
+        avg_out = torch.mean(x, dim=1, keepdim=True)
+        max_out, _ = torch.max(x, dim=1, keepdim=True)
+        x = torch.cat([avg_out, max_out], dim=1)
+        x = self.conv1(x)
+        return self.sigmoid(x)

seg2art/sstan_models/networks/base_network.py

@@ -0,0 +1,59 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import torch.nn as nn
+from torch.nn import init
+
+
+class BaseNetwork(nn.Module):
+    def __init__(self):
+        super(BaseNetwork, self).__init__()
+
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        return parser
+
+    def print_network(self):
+        if isinstance(self, list):
+            self = self[0]
+        num_params = 0
+        for param in self.parameters():
+            num_params += param.numel()
+        print('Network [%s] was created. Total number of parameters: %.1f million. '
+              'To see the architecture, do print(network).'
+              % (type(self).__name__, num_params / 1000000))
+
+    def init_weights(self, init_type='normal', gain=0.02):
+        def init_func(m):
+            classname = m.__class__.__name__
+            if classname.find('BatchNorm2d') != -1:
+                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 hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+                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 [%s] is not implemented' % init_type)
+                if hasattr(m, 'bias') and m.bias is not None:
+                    init.constant_(m.bias.data, 0.0)
+
+        self.apply(init_func)
+
+        # propagate to children
+        for m in self.children():
+            if hasattr(m, 'init_weights'):
+                m.init_weights(init_type, gain)

seg2art/sstan_models/networks/dual_attention_module.py

@@ -0,0 +1,51 @@
+
+class ChannelAttention(nn.Module):
+    def __init__(self, in_planes, ratio=16):
+        super(ChannelAttention, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.max_pool = nn.AdaptiveMaxPool2d(1)
+
+        self.fc1 = nn.Conv2d(in_planes, in_planes // 16, 1, bias=False)
+        self.relu1 = nn.ReLU()
+        self.fc2 = nn.Conv2d(in_planes // 16, in_planes, 1, bias=False)
+
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
+        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
+        out = avg_out + max_out
+        return self.sigmoid(out)
+
+
+class SpatialAttention(nn.Module):
+    def __init__(self, kernel_size=7):
+        super(SpatialAttention, self).__init__()
+
+        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
+        padding = 3 if kernel_size == 7 else 1
+
+        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = torch.mean(x, dim=1, keepdim=True)
+        max_out, _ = torch.max(x, dim=1, keepdim=True)
+        x = torch.cat([avg_out, max_out], dim=1)
+        x = self.conv1(x)
+        return self.sigmoid(x)
+
+
+'''
+    def forward(self, x, seg):
+        x_s = self.shortcut(x, seg)
+
+        dx = self.conv_0(self.actvn(self.norm_0(x, seg)))
+        dx = self.conv_1(self.actvn(self.norm_1(dx, seg)))
+
+        dx = self.channelAtt(dx) * dx
+        dx = self.spatialAtt(dx) * dx
+        out = x_s + dx
+
+        return out
+'''

seg2art/sstan_models/networks/generator.py

@@ -0,0 +1,184 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from sstan_models.networks.base_network import BaseNetwork
+from sstan_models.networks.normalization import get_nonspade_norm_layer
+from sstan_models.networks.architecture import ResnetBlock as ResnetBlock
+from sstan_models.networks.architecture import SPADEResnetBlock as SPADEResnetBlock
+import numpy as np
+torch.manual_seed(1234)
+
+
+class SPADEGenerator(BaseNetwork):
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        parser.set_defaults(norm_G='spectralspadesyncbatch3x3')
+        parser.add_argument('--num_upsampling_layers',
+                            choices=('normal', 'more', 'most'), default='normal',
+                            help="If 'more', adds upsampling layer between the two middle resnet blocks. If 'most', also add one more upsampling + resnet layer at the end of the generator")
+
+        return parser
+
+    def __init__(self, opt):
+        super().__init__()
+        self.opt = opt
+        nf = opt.ngf
+
+        self.sw, self.sh = self.compute_latent_vector_size(opt)
+
+        # if opt.use_vae:
+        #     # In case of VAE, we will sample from random z vector
+        #     self.fc = nn.Linear(opt.z_dim, 16 * nf * self.sw * self.sh)
+        # else:
+        #     # Otherwise, we make the network deterministic by starting with
+        #     # downsampled segmentation map instead of random z
+        #     self.fc = nn.Conv2d(self.opt.semantic_nc, 16 * nf, 3, padding=1)
+        self.fc = nn.Linear(opt.z_dim, 16 * nf * self.sw * self.sh)
+        self.head_0 = SPADEResnetBlock(16 * nf, 16 * nf, opt, feed_code = True)
+
+        self.G_middle_0 = SPADEResnetBlock(16 * nf, 16 * nf, opt, feed_code = True)
+        self.G_middle_1 = SPADEResnetBlock(16 * nf, 16 * nf, opt, feed_code = True)
+
+        self.up_0 = SPADEResnetBlock(16 * nf, 8 * nf, opt, feed_code = True)
+        self.up_1 = SPADEResnetBlock(8 * nf, 4 * nf, opt, feed_code = True)
+        self.up_2 = SPADEResnetBlock(4 * nf, 2 * nf, opt, feed_code = True)
+        self.up_3 = SPADEResnetBlock(2 * nf, 1 * nf, opt, feed_code = False)
+
+        final_nc = nf
+
+        # if opt.num_upsampling_layers == 'most':
+        #     print('used?')
+        #     self.up_4 = SPADEResnetBlock(1 * nf, nf // 2, opt)
+        #     final_nc = nf // 2
+
+        self.conv_img = nn.Conv2d(final_nc, 3, 3, padding=1)
+
+        self.up = nn.Upsample(scale_factor=2)
+
+    def compute_latent_vector_size(self, opt):
+        if opt.num_upsampling_layers == 'normal':
+            num_up_layers = 5
+        elif opt.num_upsampling_layers == 'more':
+            num_up_layers = 6
+        elif opt.num_upsampling_layers == 'most':
+            num_up_layers = 7
+        else:
+            raise ValueError('opt.num_upsampling_layers [%s] not recognized' %
+                             opt.num_upsampling_layers)
+
+        sw = opt.crop_size // (2**num_up_layers)
+        sh = round(sw / opt.aspect_ratio)
+
+        return sw, sh
+
+    def forward(self, input, rgb_img, style_codes=None):
+        with torch.cuda.amp.autocast():
+            seg = input
+            # if self.opt.use_vae:
+            #     # we sample z from unit normal and reshape the tensor
+            #     if z is None:
+            #         z = torch.randn(input.size(0), self.opt.z_dim,
+            #                         dtype=torch.float32, device=input.get_device())
+            x = self.fc(style_codes)
+            x = x.view(-1, 16 * self.opt.ngf, self.sh, self.sw)
+            # else:
+            #     # we downsample segmap and run convolution
+            #     x = F.interpolate(seg, size=(self.sh, self.sw))
+            #     x = self.fc(x)
+            x = self.head_0(x, seg, style_codes=style_codes)
+
+            x = self.up(x)
+            x = self.G_middle_0(x, seg, style_codes=style_codes)
+
+            if self.opt.num_upsampling_layers == 'more' or \
+            self.opt.num_upsampling_layers == 'most':
+                x = self.up(x)
+
+            x = self.G_middle_1(x, seg, style_codes=style_codes)
+
+            x = self.up(x)
+            x = self.up_0(x, seg, style_codes=style_codes)
+            x = self.up(x)
+            x = self.up_1(x, seg, style_codes=style_codes)
+            x = self.up(x)
+            x = self.up_2(x, seg, style_codes=style_codes)
+            x = self.up(x)
+            x = self.up_3(x, seg)
+
+            # if self.opt.num_upsampling_layers == 'most':
+            #     x = self.up(x)
+            #     x = self.up_4(x, seg)
+
+            x = self.conv_img(F.leaky_relu(x, 2e-1))
+            x = torch.tanh(x)#F.tanh(x)
+
+        return x#, style_codes
+
+
+class Pix2PixHDGenerator(BaseNetwork):
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        parser.add_argument('--resnet_n_downsample', type=int, default=4, help='number of downsampling layers in netG')
+        parser.add_argument('--resnet_n_blocks', type=int, default=9, help='number of residual blocks in the global generator network')
+        parser.add_argument('--resnet_kernel_size', type=int, default=3,
+                            help='kernel size of the resnet block')
+        parser.add_argument('--resnet_initial_kernel_size', type=int, default=7,
+                            help='kernel size of the first convolution')
+        parser.set_defaults(norm_G='instance')
+        return parser
+
+    def __init__(self, opt):
+        super().__init__()
+        input_nc = opt.label_nc + (1 if opt.contain_dontcare_label else 0) + (0 if opt.no_instance else 1)
+
+        norm_layer = get_nonspade_norm_layer(opt, opt.norm_G)
+        activation = nn.ReLU(False)
+
+        model = []
+
+        # initial conv
+        model += [nn.ReflectionPad2d(opt.resnet_initial_kernel_size // 2),
+                  norm_layer(nn.Conv2d(input_nc, opt.ngf,
+                                       kernel_size=opt.resnet_initial_kernel_size,
+                                       padding=0)),
+                  activation]
+
+        # downsample
+        mult = 1
+        for i in range(opt.resnet_n_downsample):
+            model += [norm_layer(nn.Conv2d(opt.ngf * mult, opt.ngf * mult * 2,
+                                           kernel_size=3, stride=2, padding=1)),
+                      activation]
+            mult *= 2
+
+        # resnet blocks
+        for i in range(opt.resnet_n_blocks):
+            model += [ResnetBlock(opt.ngf * mult,
+                                  norm_layer=norm_layer,
+                                  activation=activation,
+                                  kernel_size=opt.resnet_kernel_size)]
+
+        # upsample
+        for i in range(opt.resnet_n_downsample):
+            nc_in = int(opt.ngf * mult)
+            nc_out = int((opt.ngf * mult) / 2)
+            model += [norm_layer(nn.ConvTranspose2d(nc_in, nc_out,
+                                                    kernel_size=3, stride=2,
+                                                    padding=1, output_padding=1)),
+                      activation]
+            mult = mult // 2
+
+        # final output conv
+        model += [nn.ReflectionPad2d(3),
+                  nn.Conv2d(nc_out, opt.output_nc, kernel_size=7, padding=0),
+                  nn.Tanh()]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input, z=None):
+        return self.model(input)

seg2art/sstan_models/networks/normalization.py

@@ -0,0 +1,222 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import re
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from sstan_models.networks.sync_batchnorm import SynchronizedBatchNorm2d
+import torch.nn.utils.spectral_norm as spectral_norm
+
+
+# Returns a function that creates a normalization function
+# that does not condition on semantic map
+def get_nonspade_norm_layer(opt, norm_type='instance'):
+    # helper function to get # output channels of the previous layer
+    def get_out_channel(layer):
+        if hasattr(layer, 'out_channels'):
+            return getattr(layer, 'out_channels')
+        return layer.weight.size(0)
+
+    # this function will be returned
+    def add_norm_layer(layer):
+        nonlocal norm_type
+        if norm_type.startswith('spectral'):
+            layer = spectral_norm(layer)
+            subnorm_type = norm_type[len('spectral'):]
+
+        if subnorm_type == 'none' or len(subnorm_type) == 0:
+            return layer
+
+        # remove bias in the previous layer, which is meaningless
+        # since it has no effect after normalization
+        if getattr(layer, 'bias', None) is not None:
+            delattr(layer, 'bias')
+            layer.register_parameter('bias', None)
+
+        if subnorm_type == 'batch':
+            norm_layer = nn.BatchNorm2d(get_out_channel(layer), affine=True)
+        elif subnorm_type == 'sync_batch':
+            norm_layer = SynchronizedBatchNorm2d(
+                get_out_channel(layer), affine=True)
+        elif subnorm_type == 'instance':
+            norm_layer = nn.InstanceNorm2d(
+                get_out_channel(layer), affine=False)
+        else:
+            raise ValueError(
+                'normalization layer %s is not recognized' % subnorm_type)
+
+        return nn.Sequential(layer, norm_layer)
+
+    return add_norm_layer
+
+
+# Creates SPADE normalization layer based on the given configuration
+# SPADE consists of two steps. First, it normalizes the activations using
+# your favorite normalization method, such as Batch Norm or Instance Norm.
+# Second, it applies scale and bias to the normalized output, conditioned on
+# the segmentation map.
+# The format of |config_text| is spade(norm)(ks), where
+# (norm) specifies the type of parameter-free normalization.
+#       (e.g. syncbatch, batch, instance)
+# (ks) specifies the size of kernel in the SPADE module (e.g. 3x3)
+# Example |config_text| will be spadesyncbatch3x3, or spadeinstance5x5.
+# Also, the other arguments are
+# |norm_nc|: the #channels of the normalized activations, hence the output dim of SPADE
+# |label_nc|: the #channels of the input semantic map, hence the input dim of SPADE
+class SPADE(nn.Module):
+    def __init__(self, config_text, norm_nc, feed_code, status='train', spade_params=None):
+        super().__init__()
+
+        self.style_length = 256
+        # self.noise_var = nn.Parameter(torch.zeros(norm_nc), requires_grad=True)
+        self.Spade = SPADE_ori(*spade_params)
+
+
+        assert config_text.startswith('spade')
+        parsed = re.search('spade(\D+)(\d)x\d', config_text)
+        param_free_norm_type = str(parsed.group(1))
+        ks = int(parsed.group(2))
+        pw = ks // 2
+
+        if param_free_norm_type == 'instance':
+            self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)
+        elif param_free_norm_type == 'syncbatch':
+            self.param_free_norm = SynchronizedBatchNorm2d(
+                norm_nc, affine=False)
+        elif param_free_norm_type == 'batch':
+            self.param_free_norm = nn.BatchNorm2d(norm_nc, affine=False)
+        else:
+            raise ValueError('%s is not a recognized param-free norm type in SPADE'
+                             % param_free_norm_type)
+
+        # self.create_gamma_beta_fc_layers()
+        if feed_code:
+            self.blending_gamma = nn.Parameter(torch.zeros(1), requires_grad=True)
+            self.blending_beta = nn.Parameter(torch.zeros(1), requires_grad=True)
+            self.conv_gamma = nn.Conv2d(
+                self.style_length, norm_nc, kernel_size=ks, padding=pw)
+            self.conv_beta = nn.Conv2d(
+                self.style_length, norm_nc, kernel_size=ks, padding=pw)
+
+    def forward(self, x, segmap, style_codes=None):
+        if style_codes is None:
+            input_code = False
+        else:
+            input_code = True
+
+        # Part 1. generate parameter-free normalized activations
+        # added_noise = (torch.randn(
+        #     x.shape[0], x.shape[3], x.shape[2], 1).cuda() * self.noise_var).transpose(1, 3)
+        normalized = self.param_free_norm(x)
+
+        # Part 2. produce scaling and bias conditioned on semantic map
+        segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest')
+
+        if input_code:
+            [b_size, f_size, h_size, w_size] = normalized.shape
+            middle_avg = torch.zeros(
+                (b_size, self.style_length, h_size, w_size), device=normalized.device)
+
+            for i in range(b_size):
+
+                middle_mu = F.relu((style_codes[i]))
+
+                middle_mu = middle_mu.reshape(self.style_length, 1).expand(
+                    self.style_length, h_size*w_size)
+                middle_mu = middle_mu.reshape(
+                    self.style_length, h_size, w_size)
+                middle_avg[i] = middle_mu
+
+            gamma_avg = self.conv_gamma(middle_avg)
+            beta_avg = self.conv_beta(middle_avg)
+
+            gamma_spade, beta_spade = self.Spade(segmap)
+
+            gamma_alpha = torch.sigmoid(self.blending_gamma)#F.sigmoid(self.blending_gamma)
+            beta_alpha = torch.sigmoid(self.blending_gamma)#F.sigmoid(self.blending_beta)
+
+            gamma_final = gamma_alpha * gamma_avg + \
+                (1 - gamma_alpha) * gamma_spade
+
+            beta_final = beta_alpha * beta_avg + (1 - beta_alpha) * beta_spade
+
+            out = normalized * (1 + gamma_final) + beta_final
+        else:
+            gamma_spade, beta_spade = self.Spade(segmap)
+            gamma_final = gamma_spade
+            beta_final = beta_spade
+            out = normalized * (1 + gamma_final) + beta_final
+        return out
+
+    # def create_gamma_beta_fc_layers(self):
+
+    #     # These codes should be replaced with torch.nn.ModuleList
+
+    #     style_length = self.style_length
+
+    #     self.fc_mu0 = nn.Linear(style_length, style_length)
+    #     self.fc_mu1 = nn.Linear(style_length, style_length)
+    #     self.fc_mu2 = nn.Linear(style_length, style_length)
+    #     self.fc_mu3 = nn.Linear(style_length, style_length)
+    #     self.fc_mu4 = nn.Linear(style_length, style_length)
+    #     self.fc_mu5 = nn.Linear(style_length, style_length)
+    #     self.fc_mu6 = nn.Linear(style_length, style_length)
+    #     self.fc_mu7 = nn.Linear(style_length, style_length)
+    #     self.fc_mu8 = nn.Linear(style_length, style_length)
+    #     self.fc_mu9 = nn.Linear(style_length, style_length)
+    #     self.fc_mu10 = nn.Linear(style_length, style_length)
+    #     self.fc_mu11 = nn.Linear(style_length, style_length)
+    #     self.fc_mu12 = nn.Linear(style_length, style_length)
+    #     self.fc_mu13 = nn.Linear(style_length, style_length)
+    #     self.fc_mu14 = nn.Linear(style_length, style_length)
+    #     self.fc_mu15 = nn.Linear(style_length, style_length)
+    #     self.fc_mu16 = nn.Linear(style_length, style_length)
+    #     self.fc_mu17 = nn.Linear(style_length, style_length)
+    #     self.fc_mu18 = nn.Linear(style_length, style_length)
+
+
+class SPADE_ori(nn.Module):
+    def __init__(self, config_text, norm_nc, label_nc):
+        super().__init__()
+
+        assert config_text.startswith('spade')
+        parsed = re.search('spade(\D+)(\d)x\d', config_text)
+        param_free_norm_type = str(parsed.group(1))
+        ks = int(parsed.group(2))
+
+        if param_free_norm_type == 'instance':
+            self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)
+        elif param_free_norm_type == 'syncbatch':
+            self.param_free_norm = SynchronizedBatchNorm2d(
+                norm_nc, affine=False)
+        elif param_free_norm_type == 'batch':
+            self.param_free_norm = nn.BatchNorm2d(norm_nc, affine=False)
+        else:
+            raise ValueError('%s is not a recognized param-free norm type in SPADE'
+                             % param_free_norm_type)
+
+        # The dimension of the intermediate embedding space. Yes, hardcoded.
+        nhidden = 128
+
+        pw = ks // 2
+        self.mlp_shared = nn.Sequential(
+            nn.Conv2d(label_nc, nhidden, kernel_size=ks, padding=pw),
+            nn.ReLU()
+        )
+
+        self.mlp_gamma = nn.Conv2d(
+            nhidden, norm_nc, kernel_size=ks, padding=pw)
+        self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw)
+
+    def forward(self, segmap):
+
+        inputmap = segmap
+
+        actv = self.mlp_shared(inputmap)
+        gamma = self.mlp_gamma(actv)
+        beta = self.mlp_beta(actv)
+
+        return gamma, beta

seg2art/sstan_models/networks/sync_batchnorm/__init__.py

@@ -0,0 +1,13 @@
+# -*- coding: utf-8 -*-
+# File   : __init__.py
+# Author : Jiayuan Mao
+# Email  : [email protected]
+# Date   : 27/01/2018
+# 
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+from .batchnorm import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, SynchronizedBatchNorm3d
+from .batchnorm import convert_model
+from .replicate import DataParallelWithCallback, patch_replication_callback

seg2art/sstan_models/networks/sync_batchnorm/batchnorm.py

@@ -0,0 +1,361 @@
+# -*- coding: utf-8 -*-
+# File   : batchnorm.py
+# Author : Jiayuan Mao
+# Email  : [email protected]
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import collections
+
+import torch
+import torch.nn.functional as F
+
+from torch.nn.modules.batchnorm import _BatchNorm
+from torch.nn.parallel._functions import ReduceAddCoalesced, Broadcast
+
+from .comm import SyncMaster
+from .replicate import DataParallelWithCallback
+
+__all__ = ['SynchronizedBatchNorm1d', 'SynchronizedBatchNorm2d',
+           'SynchronizedBatchNorm3d', 'convert_model']
+
+
+def _sum_ft(tensor):
+    """sum over the first and last dimention"""
+    return tensor.sum(dim=0).sum(dim=-1)
+
+
+def _unsqueeze_ft(tensor):
+    """add new dementions at the front and the tail"""
+    return tensor.unsqueeze(0).unsqueeze(-1)
+
+
+_ChildMessage = collections.namedtuple('_ChildMessage', ['sum', 'ssum', 'sum_size'])
+_MasterMessage = collections.namedtuple('_MasterMessage', ['sum', 'inv_std'])
+
+
+class _SynchronizedBatchNorm(_BatchNorm):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True):
+        super(_SynchronizedBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine)
+
+        self._sync_master = SyncMaster(self._data_parallel_master)
+
+        self._is_parallel = False
+        self._parallel_id = None
+        self._slave_pipe = None
+
+    def forward(self, input):
+        # If it is not parallel computation or is in evaluation mode, use PyTorch's implementation.
+        if not (self._is_parallel and self.training):
+            return F.batch_norm(
+                input, self.running_mean, self.running_var, self.weight, self.bias,
+                self.training, self.momentum, self.eps)
+
+        # Resize the input to (B, C, -1).
+        input_shape = input.size()
+        input = input.view(input.size(0), self.num_features, -1)
+
+        # Compute the sum and square-sum.
+        sum_size = input.size(0) * input.size(2)
+        input_sum = _sum_ft(input)
+        input_ssum = _sum_ft(input ** 2)
+
+        # Reduce-and-broadcast the statistics.
+        if self._parallel_id == 0:
+            mean, inv_std = self._sync_master.run_master(_ChildMessage(input_sum, input_ssum, sum_size))
+        else:
+            mean, inv_std = self._slave_pipe.run_slave(_ChildMessage(input_sum, input_ssum, sum_size))
+
+        # Compute the output.
+        if self.affine:
+            # MJY:: Fuse the multiplication for speed.
+            output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std * self.weight) + _unsqueeze_ft(self.bias)
+        else:
+            output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std)
+
+        # Reshape it.
+        return output.view(input_shape)
+
+    def __data_parallel_replicate__(self, ctx, copy_id):
+        self._is_parallel = True
+        self._parallel_id = copy_id
+
+        # parallel_id == 0 means master device.
+        if self._parallel_id == 0:
+            ctx.sync_master = self._sync_master
+        else:
+            self._slave_pipe = ctx.sync_master.register_slave(copy_id)
+
+    def _data_parallel_master(self, intermediates):
+        """Reduce the sum and square-sum, compute the statistics, and broadcast it."""
+
+        # Always using same "device order" makes the ReduceAdd operation faster.
+        # Thanks to:: Tete Xiao (http://tetexiao.com/)
+        intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device())
+
+        to_reduce = [i[1][:2] for i in intermediates]
+        to_reduce = [j for i in to_reduce for j in i]  # flatten
+        target_gpus = [i[1].sum.get_device() for i in intermediates]
+
+        sum_size = sum([i[1].sum_size for i in intermediates])
+        sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce)
+        mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size)
+
+        broadcasted = Broadcast.apply(target_gpus, mean, inv_std)
+
+        outputs = []
+        for i, rec in enumerate(intermediates):
+            outputs.append((rec[0], _MasterMessage(*broadcasted[i*2:i*2+2])))
+
+        return outputs
+
+    def _compute_mean_std(self, sum_, ssum, size):
+        """Compute the mean and standard-deviation with sum and square-sum. This method
+        also maintains the moving average on the master device."""
+        assert size > 1, 'BatchNorm computes unbiased standard-deviation, which requires size > 1.'
+        mean = sum_ / size
+        sumvar = ssum - sum_ * mean
+        unbias_var = sumvar / (size - 1)
+        bias_var = sumvar / size
+
+        self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data
+        self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data
+
+        return mean, bias_var.clamp(self.eps) ** -0.5
+
+
+class SynchronizedBatchNorm1d(_SynchronizedBatchNorm):
+    r"""Applies Synchronized Batch Normalization over a 2d or 3d input that is seen as a
+    mini-batch.
+
+    .. math::
+
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+
+    This module differs from the built-in PyTorch BatchNorm1d as the mean and
+    standard-deviation are reduced across all devices during training.
+
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+
+    During evaluation, this running mean/variance is used for normalization.
+
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, L)` slices, it's common terminology to call this Temporal BatchNorm
+
+    Args:
+        num_features: num_features from an expected input of size
+            `batch_size x num_features [x width]`
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C)` or :math:`(N, C, L)`
+        - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input)
+
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm1d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm1d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 2 and input.dim() != 3:
+            raise ValueError('expected 2D or 3D input (got {}D input)'
+                             .format(input.dim()))
+        super(SynchronizedBatchNorm1d, self)._check_input_dim(input)
+
+
+class SynchronizedBatchNorm2d(_SynchronizedBatchNorm):
+    r"""Applies Batch Normalization over a 4d input that is seen as a mini-batch
+    of 3d inputs
+
+    .. math::
+
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+
+    This module differs from the built-in PyTorch BatchNorm2d as the mean and
+    standard-deviation are reduced across all devices during training.
+
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+
+    During evaluation, this running mean/variance is used for normalization.
+
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, H, W)` slices, it's common terminology to call this Spatial BatchNorm
+
+    Args:
+        num_features: num_features from an expected input of
+            size batch_size x num_features x height x width
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C, H, W)`
+        - Output: :math:`(N, C, H, W)` (same shape as input)
+
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm2d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm2d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 4:
+            raise ValueError('expected 4D input (got {}D input)'
+                             .format(input.dim()))
+        super(SynchronizedBatchNorm2d, self)._check_input_dim(input)
+
+
+class SynchronizedBatchNorm3d(_SynchronizedBatchNorm):
+    r"""Applies Batch Normalization over a 5d input that is seen as a mini-batch
+    of 4d inputs
+
+    .. math::
+
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+
+    This module differs from the built-in PyTorch BatchNorm3d as the mean and
+    standard-deviation are reduced across all devices during training.
+
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+
+    During evaluation, this running mean/variance is used for normalization.
+
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, D, H, W)` slices, it's common terminology to call this Volumetric BatchNorm
+    or Spatio-temporal BatchNorm
+
+    Args:
+        num_features: num_features from an expected input of
+            size batch_size x num_features x depth x height x width
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+
+    Shape:
+        - Input: :math:`(N, C, D, H, W)`
+        - Output: :math:`(N, C, D, H, W)` (same shape as input)
+
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm3d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm3d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45, 10))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 5:
+            raise ValueError('expected 5D input (got {}D input)'
+                             .format(input.dim()))
+        super(SynchronizedBatchNorm3d, self)._check_input_dim(input)
+
+        
+def convert_model(module):
+    """Traverse the input module and its child recursively
+       and replace all instance of torch.nn.modules.batchnorm.BatchNorm*N*d
+       to SynchronizedBatchNorm*N*d 
+        
+    Args:
+        module: the input module needs to be convert to SyncBN model
+        
+    Examples:
+        >>> import torch.nn as nn
+        >>> import torchvision
+        >>> # m is a standard pytorch model
+        >>> m = torchvision.models.resnet18(True)
+        >>> m = nn.DataParallel(m)
+        >>> # after convert, m is using SyncBN
+        >>> m = convert_model(m)
+    """
+    if isinstance(module, torch.nn.DataParallel):
+        mod = module.module
+        mod = convert_model(mod)
+        mod = DataParallelWithCallback(mod)
+        return mod
+    
+    mod = module
+    for pth_module, sync_module in zip([torch.nn.modules.batchnorm.BatchNorm1d,
+                                        torch.nn.modules.batchnorm.BatchNorm2d,
+                                        torch.nn.modules.batchnorm.BatchNorm3d],
+                                       [SynchronizedBatchNorm1d,
+                                        SynchronizedBatchNorm2d,
+                                        SynchronizedBatchNorm3d]):
+        if isinstance(module, pth_module):
+            mod = sync_module(module.num_features, module.eps, module.momentum, module.affine)
+            mod.running_mean = module.running_mean
+            mod.running_var = module.running_var
+            if module.affine:
+                mod.weight.data = module.weight.data.clone().detach()
+                mod.bias.data = module.bias.data.clone().detach()
+            
+    for name, child in module.named_children():
+        mod.add_module(name, convert_model(child))
+    
+    return mod

seg2art/sstan_models/networks/sync_batchnorm/batchnorm_reimpl.py

@@ -0,0 +1,74 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+# File   : batchnorm_reimpl.py
+# Author : acgtyrant
+# Date   : 11/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+
+__all__ = ['BatchNormReimpl']
+
+
+class BatchNorm2dReimpl(nn.Module):
+    """
+    A re-implementation of batch normalization, used for testing the numerical
+    stability.
+
+    Author: acgtyrant
+    See also:
+    https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues/14
+    """
+    def __init__(self, num_features, eps=1e-5, momentum=0.1):
+        super().__init__()
+
+        self.num_features = num_features
+        self.eps = eps
+        self.momentum = momentum
+        self.weight = nn.Parameter(torch.empty(num_features))
+        self.bias = nn.Parameter(torch.empty(num_features))
+        self.register_buffer('running_mean', torch.zeros(num_features))
+        self.register_buffer('running_var', torch.ones(num_features))
+        self.reset_parameters()
+
+    def reset_running_stats(self):
+        self.running_mean.zero_()
+        self.running_var.fill_(1)
+
+    def reset_parameters(self):
+        self.reset_running_stats()
+        init.uniform_(self.weight)
+        init.zeros_(self.bias)
+
+    def forward(self, input_):
+        batchsize, channels, height, width = input_.size()
+        numel = batchsize * height * width
+        input_ = input_.permute(1, 0, 2, 3).contiguous().view(channels, numel)
+        sum_ = input_.sum(1)
+        sum_of_square = input_.pow(2).sum(1)
+        mean = sum_ / numel
+        sumvar = sum_of_square - sum_ * mean
+
+        self.running_mean = (
+                (1 - self.momentum) * self.running_mean
+                + self.momentum * mean.detach()
+        )
+        unbias_var = sumvar / (numel - 1)
+        self.running_var = (
+                (1 - self.momentum) * self.running_var
+                + self.momentum * unbias_var.detach()
+        )
+
+        bias_var = sumvar / numel
+        inv_std = 1 / (bias_var + self.eps).pow(0.5)
+        output = (
+                (input_ - mean.unsqueeze(1)) * inv_std.unsqueeze(1) *
+                self.weight.unsqueeze(1) + self.bias.unsqueeze(1))
+
+        return output.view(channels, batchsize, height, width).permute(1, 0, 2, 3).contiguous()
+

seg2art/sstan_models/networks/sync_batchnorm/comm.py

@@ -0,0 +1,137 @@
+# -*- coding: utf-8 -*-
+# File   : comm.py
+# Author : Jiayuan Mao
+# Email  : [email protected]
+# Date   : 27/01/2018
+# 
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import queue
+import collections
+import threading
+
+__all__ = ['FutureResult', 'SlavePipe', 'SyncMaster']
+
+
+class FutureResult(object):
+    """A thread-safe future implementation. Used only as one-to-one pipe."""
+
+    def __init__(self):
+        self._result = None
+        self._lock = threading.Lock()
+        self._cond = threading.Condition(self._lock)
+
+    def put(self, result):
+        with self._lock:
+            assert self._result is None, 'Previous result has\'t been fetched.'
+            self._result = result
+            self._cond.notify()
+
+    def get(self):
+        with self._lock:
+            if self._result is None:
+                self._cond.wait()
+
+            res = self._result
+            self._result = None
+            return res
+
+
+_MasterRegistry = collections.namedtuple('MasterRegistry', ['result'])
+_SlavePipeBase = collections.namedtuple('_SlavePipeBase', ['identifier', 'queue', 'result'])
+
+
+class SlavePipe(_SlavePipeBase):
+    """Pipe for master-slave communication."""
+
+    def run_slave(self, msg):
+        self.queue.put((self.identifier, msg))
+        ret = self.result.get()
+        self.queue.put(True)
+        return ret
+
+
+class SyncMaster(object):
+    """An abstract `SyncMaster` object.
+
+    - During the replication, as the data parallel will trigger an callback of each module, all slave devices should
+    call `register(id)` and obtain an `SlavePipe` to communicate with the master.
+    - During the forward pass, master device invokes `run_master`, all messages from slave devices will be collected,
+    and passed to a registered callback.
+    - After receiving the messages, the master device should gather the information and determine to message passed
+    back to each slave devices.
+    """
+
+    def __init__(self, master_callback):
+        """
+
+        Args:
+            master_callback: a callback to be invoked after having collected messages from slave devices.
+        """
+        self._master_callback = master_callback
+        self._queue = queue.Queue()
+        self._registry = collections.OrderedDict()
+        self._activated = False
+
+    def __getstate__(self):
+        return {'master_callback': self._master_callback}
+
+    def __setstate__(self, state):
+        self.__init__(state['master_callback'])
+
+    def register_slave(self, identifier):
+        """
+        Register an slave device.
+
+        Args:
+            identifier: an identifier, usually is the device id.
+
+        Returns: a `SlavePipe` object which can be used to communicate with the master device.
+
+        """
+        if self._activated:
+            assert self._queue.empty(), 'Queue is not clean before next initialization.'
+            self._activated = False
+            self._registry.clear()
+        future = FutureResult()
+        self._registry[identifier] = _MasterRegistry(future)
+        return SlavePipe(identifier, self._queue, future)
+
+    def run_master(self, master_msg):
+        """
+        Main entry for the master device in each forward pass.
+        The messages were first collected from each devices (including the master device), and then
+        an callback will be invoked to compute the message to be sent back to each devices
+        (including the master device).
+
+        Args:
+            master_msg: the message that the master want to send to itself. This will be placed as the first
+            message when calling `master_callback`. For detailed usage, see `_SynchronizedBatchNorm` for an example.
+
+        Returns: the message to be sent back to the master device.
+
+        """
+        self._activated = True
+
+        intermediates = [(0, master_msg)]
+        for i in range(self.nr_slaves):
+            intermediates.append(self._queue.get())
+
+        results = self._master_callback(intermediates)
+        assert results[0][0] == 0, 'The first result should belongs to the master.'
+
+        for i, res in results:
+            if i == 0:
+                continue
+            self._registry[i].result.put(res)
+
+        for i in range(self.nr_slaves):
+            assert self._queue.get() is True
+
+        return results[0][1]
+
+    @property
+    def nr_slaves(self):
+        return len(self._registry)

seg2art/sstan_models/networks/sync_batchnorm/replicate.py

@@ -0,0 +1,94 @@
+# -*- coding: utf-8 -*-
+# File   : replicate.py
+# Author : Jiayuan Mao
+# Email  : [email protected]
+# Date   : 27/01/2018
+# 
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import functools
+
+from torch.nn.parallel.data_parallel import DataParallel
+
+__all__ = [
+    'CallbackContext',
+    'execute_replication_callbacks',
+    'DataParallelWithCallback',
+    'patch_replication_callback'
+]
+
+
+class CallbackContext(object):
+    pass
+
+
+def execute_replication_callbacks(modules):
+    """
+    Execute an replication callback `__data_parallel_replicate__` on each module created by original replication.
+
+    The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)`
+
+    Note that, as all modules are isomorphism, we assign each sub-module with a context
+    (shared among multiple copies of this module on different devices).
+    Through this context, different copies can share some information.
+
+    We guarantee that the callback on the master copy (the first copy) will be called ahead of calling the callback
+    of any slave copies.
+    """
+    master_copy = modules[0]
+    nr_modules = len(list(master_copy.modules()))
+    ctxs = [CallbackContext() for _ in range(nr_modules)]
+
+    for i, module in enumerate(modules):
+        for j, m in enumerate(module.modules()):
+            if hasattr(m, '__data_parallel_replicate__'):
+                m.__data_parallel_replicate__(ctxs[j], i)
+
+
+class DataParallelWithCallback(DataParallel):
+    """
+    Data Parallel with a replication callback.
+
+    An replication callback `__data_parallel_replicate__` of each module will be invoked after being created by
+    original `replicate` function.
+    The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)`
+
+    Examples:
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+        # sync_bn.__data_parallel_replicate__ will be invoked.
+    """
+
+    def replicate(self, module, device_ids):
+        modules = super(DataParallelWithCallback, self).replicate(module, device_ids)
+        execute_replication_callbacks(modules)
+        return modules
+
+
+def patch_replication_callback(data_parallel):
+    """
+    Monkey-patch an existing `DataParallel` object. Add the replication callback.
+    Useful when you have customized `DataParallel` implementation.
+
+    Examples:
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallel(sync_bn, device_ids=[0, 1])
+        > patch_replication_callback(sync_bn)
+        # this is equivalent to
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+    """
+
+    assert isinstance(data_parallel, DataParallel)
+
+    old_replicate = data_parallel.replicate
+
+    @functools.wraps(old_replicate)
+    def new_replicate(module, device_ids):
+        modules = old_replicate(module, device_ids)
+        execute_replication_callbacks(modules)
+        return modules
+
+    data_parallel.replicate = new_replicate

seg2art/sstan_models/networks/sync_batchnorm/unittest.py

@@ -0,0 +1,29 @@
+# -*- coding: utf-8 -*-
+# File   : unittest.py
+# Author : Jiayuan Mao
+# Email  : [email protected]
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import unittest
+import torch
+
+
+class TorchTestCase(unittest.TestCase):
+    def assertTensorClose(self, x, y):
+        adiff = float((x - y).abs().max())
+        if (y == 0).all():
+            rdiff = 'NaN'
+        else:
+            rdiff = float((adiff / y).abs().max())
+
+        message = (
+            'Tensor close check failed\n'
+            'adiff={}\n'
+            'rdiff={}\n'
+        ).format(adiff, rdiff)
+        self.assertTrue(torch.allclose(x, y), message)
+

seg2art/sstan_models/pix2pix_model.py

@@ -0,0 +1,285 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import os
+import torch
+import sstan_models.networks as networks
+import model_util as util
+
+
+class Pix2PixModel(torch.nn.Module):
+    @staticmethod
+    def modify_commandline_options(parser, is_train):
+        networks.modify_commandline_options(parser, is_train)
+        return parser
+
+    def __init__(self, opt):
+        super().__init__()
+        self.opt = opt
+        self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() else torch.FloatTensor
+        self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() else torch.ByteTensor
+
+        self.netG, self.netD, self.netE = self.initialize_networks(opt)
+
+        # set loss functions
+        if opt.isTrain:
+            self.criterionGAN = networks.GANLoss(opt.gan_mode, tensor=self.FloatTensor, opt=self.opt)
+            self.criterionFeat = torch.nn.L1Loss()
+            if not opt.no_vgg_loss:
+                self.criterionVGG = networks.VGGLoss(self.opt.gpu_ids)
+            if opt.use_vae:
+                self.KLDLoss = networks.KLDLoss()
+
+    # Entry point for all calls involving forward pass
+    # of deep networks. We used this approach since DataParallel module
+    # can't parallelize custom functions, we branch to different
+    # routines based on |mode|.
+    def forward(self, data, mode, style_codes=None):
+        input_semantics, real_image = self.preprocess_input(data)
+        domain = None
+
+        # print(torch.cuda.memory_cached(0))
+        if mode == "generator":
+            g_loss, generated = self.compute_generator_loss(input_semantics, real_image, domain)
+            return g_loss, generated
+        elif mode == "discriminator":
+            d_loss = self.compute_discriminator_loss(input_semantics, real_image, domain)
+            return d_loss
+        elif mode == "encode_only":
+            _, mu, logvar = self.encode_z(real_image, domain)
+            return mu, logvar
+        elif mode == "inference":
+            with torch.no_grad():
+                fake_image, _, _ = self.generate_fake(input_semantics, real_image, domain, style_codes=style_codes, compute_kld_loss=False)
+            return fake_image
+        elif mode == "generate_img_npy":
+            with torch.no_grad():
+                fake_image, encoded_style_code = self.generate_img_npy(input_semantics, real_image, domain)
+            return fake_image, encoded_style_code
+        else:
+            raise ValueError("|mode| is invalid")
+
+    def create_optimizers(self, opt):
+        G_params = list(self.netG.parameters())
+        if opt.use_vae:
+            G_params += list(self.netE.parameters())
+        if opt.isTrain:
+            D_params = list(self.netD.parameters())
+
+        beta1, beta2 = opt.beta1, opt.beta2
+        if opt.no_TTUR:
+            G_lr, D_lr = opt.lr, opt.lr
+        else:
+            G_lr, D_lr = opt.lr / 2, opt.lr * 2
+
+        optimizer_G = torch.optim.Adam(G_params, lr=G_lr, betas=(beta1, beta2))
+        optimizer_D = torch.optim.Adam(D_params, lr=D_lr, betas=(beta1, beta2))
+
+        return optimizer_G, optimizer_D
+
+    def save(self, epoch):
+        util.save_network(self.netG, "G", epoch, self.opt)
+        util.save_network(self.netD, "D", epoch, self.opt)
+        if self.opt.use_vae:
+            util.save_network(self.netE, "E", epoch, self.opt)
+
+    ############################################################################
+    # Private helper methods
+    ############################################################################
+
+    def initialize_networks(self, opt):
+        netG = networks.define_G(opt)
+        netD = networks.define_D(opt) if opt.isTrain else None
+        netE = networks.define_E(opt) if opt.use_vae else None
+
+        if not opt.isTrain or opt.continue_train:
+            # netG = util.load_network(netG, 'G', opt.which_epoch, opt)
+            checkpoint_path = os.path.join(os.path.dirname(__file__), "..", opt.checkpoint_path)
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            if device == "cuda":
+                checkpoint = torch.load(checkpoint_path)
+            else:
+                checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
+            s = checkpoint
+            netG.load_state_dict(s)
+
+            if opt.isTrain:
+                netD = util.load_network(netD, "D", opt.which_epoch, opt)
+            if opt.use_vae:
+                netE = util.load_network(netE, "E", opt.which_epoch, opt)
+
+        return netG, netD, netE
+
+    # preprocess the input, such as moving the tensors to GPUs and
+    # transforming the label map to one-hot encoding
+    # |data|: dictionary of the input data
+
+    def preprocess_input(self, data):
+        """
+        # move to GPU and change data types
+        data['label'] = data['label'].long()
+        if self.use_gpu():
+            data['label'] = data['label'].cuda(non_blocking=True)
+            data['instance'] = data['instance'].cuda(non_blocking=True)
+            data['image'] = data['image'].cuda(non_blocking=True)
+            data['domain'] = data['domain'].cuda(non_blocking=True)
+
+        # create one-hot label map
+        label_map = data['label']
+        bs, _, h, w = label_map.size()
+        nc = self.opt.label_nc + 1 if self.opt.contain_dontcare_label \
+            else self.opt.label_nc
+        input_label = self.FloatTensor(bs, nc, h, w).zero_()
+        input_semantics = input_label.scatter_(1, label_map, 1.0)
+
+        # concatenate instance map if it exists
+        if not self.opt.no_instance:
+            inst_map = data['instance']
+            instance_edge_map = self.get_edges(inst_map)
+            input_semantics = torch.cat((input_semantics, instance_edge_map), dim=1)
+
+        return input_semantics, data['image'], data['domain']
+        """
+
+        data = data.long()
+        image = (data - 128).float() / 128.0
+        if self.use_gpu():
+            data = data.cuda()
+            image = image.float().cuda()
+        label_map = data
+        bs, _, h, w = label_map.size()
+        nc = self.opt.label_nc + 1 if self.opt.contain_dontcare_label else self.opt.label_nc
+        input_label = self.FloatTensor(bs, nc, h, w).zero_()
+        input_semantics = input_label.scatter_(1, label_map, 1.0)
+
+        return input_semantics, image  # data['image'],
+
+    def compute_generator_loss(self, input_semantics, real_image, domain):
+        G_losses = {}
+
+        fake_image, KLD_loss, _ = self.generate_fake(input_semantics, real_image, domain, compute_kld_loss=self.opt.use_vae)
+
+        if self.opt.use_vae:
+            if KLD_loss.data.item() > 2.5:
+                print("ng")
+                print(KLD_loss.data.item())
+            KLD_loss.data = torch.Tensor([min(999.9999, KLD_loss.data.item())]).cuda()
+            G_losses["KLD"] = KLD_loss
+
+        pred_fake, pred_real = self.discriminate(input_semantics, fake_image, real_image, domain)
+
+        G_losses["GAN"] = self.criterionGAN(pred_fake, True, for_discriminator=False)
+
+        if not self.opt.no_ganFeat_loss:
+            num_D = len(pred_fake)
+            GAN_Feat_loss = self.FloatTensor(1).fill_(0)
+            for i in range(num_D):  # for each discriminator
+                # last output is the final prediction, so we exclude it
+                num_intermediate_outputs = len(pred_fake[i]) - 1
+                for j in range(num_intermediate_outputs):  # for each layer output
+                    unweighted_loss = self.criterionFeat(pred_fake[i][j], pred_real[i][j].detach())
+                    GAN_Feat_loss += unweighted_loss * self.opt.lambda_feat / num_D
+            G_losses["GAN_Feat"] = GAN_Feat_loss
+
+        if not self.opt.no_vgg_loss:
+            G_losses["VGG"] = self.criterionVGG(fake_image, real_image) * self.opt.lambda_vgg
+
+        return G_losses, fake_image
+
+    def compute_discriminator_loss(self, input_semantics, real_image, domain):
+        D_losses = {}
+        with torch.no_grad():
+            fake_image, _, _ = self.generate_fake(input_semantics, real_image, domain)
+            fake_image = fake_image.detach()
+            fake_image.requires_grad_()
+
+        pred_fake, pred_real = self.discriminate(input_semantics, fake_image, real_image, domain)
+
+        D_losses["D_Fake"] = self.criterionGAN(pred_fake, False, for_discriminator=True)
+        D_losses["D_real"] = self.criterionGAN(pred_real, True, for_discriminator=True)
+
+        return D_losses
+
+    def encode_z(self, real_image, domain):
+        mu, logvar = self.netE(real_image, domain)
+        z = self.reparameterize(mu, logvar)
+        return z, mu, logvar
+
+    def generate_fake(self, input_semantics, real_image, domain, style_codes=None, compute_kld_loss=True):
+        KLD_loss = None
+        if self.opt.use_vae and style_codes is None:
+            # print('yes')
+            style_codes, mu, logvar = self.encode_z(real_image, domain)
+            if compute_kld_loss:
+                KLD_loss = self.KLDLoss(mu, logvar) * self.opt.lambda_kld
+
+        fake_image = self.netG(input_semantics, real_image, style_codes=style_codes)
+
+        assert (not compute_kld_loss) or self.opt.use_vae, "You cannot compute KLD loss if opt.use_vae == False"
+
+        return fake_image, KLD_loss, style_codes
+
+    def generate_img_npy(self, input_semantics, real_image, domain, compute_kld_loss=False):
+        KLD_loss = None
+        style_codes, mu, logvar = self.encode_z(real_image, domain)
+        if compute_kld_loss:
+            KLD_loss = self.KLDLoss(mu, logvar) * self.opt.lambda_kld
+
+        fake_image = self.netG(input_semantics, real_image, style_codes=style_codes)
+        # print(real_image, fake_image.shape)
+        assert (not compute_kld_loss) or self.opt.use_vae, "You cannot compute KLD loss if opt.use_vae == False"
+
+        return fake_image, style_codes
+
+    # Given fake and real image, return the prediction of discriminator
+    # for each fake and real image.
+
+    def discriminate(self, input_semantics, fake_image, real_image, domain):
+        fake_concat = torch.cat([input_semantics, fake_image], dim=1)
+        real_concat = torch.cat([input_semantics, real_image], dim=1)
+
+        # In Batch Normalization, the fake and real images are
+        # recommended to be in the same batch to avoid disparate
+        # statistics in fake and real images.
+        # So both fake and real images are fed to D all at once.
+        fake_and_real = torch.cat([fake_concat, real_concat], dim=0)
+
+        discriminator_out = self.netD(fake_and_real, domain)
+
+        pred_fake, pred_real = self.divide_pred(discriminator_out)
+
+        return pred_fake, pred_real
+
+    # Take the prediction of fake and real images from the combined batch
+    def divide_pred(self, pred):
+        # the prediction contains the intermediate outputs of multiscale GAN,
+        # so it's usually a list
+        if type(pred) == list:
+            fake = []
+            real = []
+            for p in pred:
+                fake.append([tensor[: tensor.size(0) // 2] for tensor in p])
+                real.append([tensor[tensor.size(0) // 2 :] for tensor in p])
+        else:
+            fake = pred[: pred.size(0) // 2]
+            real = pred[pred.size(0) // 2 :]
+
+        return fake, real
+
+    def get_edges(self, t):
+        edge = self.ByteTensor(t.size()).zero_()
+        edge[:, :, :, 1:] = edge[:, :, :, 1:] | (t[:, :, :, 1:] != t[:, :, :, :-1])
+        edge[:, :, :, :-1] = edge[:, :, :, :-1] | (t[:, :, :, 1:] != t[:, :, :, :-1])
+        edge[:, :, 1:, :] = edge[:, :, 1:, :] | (t[:, :, 1:, :] != t[:, :, :-1, :])
+        edge[:, :, :-1, :] = edge[:, :, :-1, :] | (t[:, :, 1:, :] != t[:, :, :-1, :])
+        return edge.float()
+
+    def reparameterize(self, mu, logvar):
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        return eps.mul(std) + mu
+
+    def use_gpu(self):
+        return len(self.opt.gpu_ids) > 0

static/index.js

@@ -0,0 +1,256 @@
+
+
+let cvsIn = document.getElementById("inputimg");
+let ctxIn = cvsIn.getContext('2d');
+let style = document.getElementById("style");
+let svgGraph = null;
+let mouselbtn = false;
+var current_time = (new Date()).getTime();
+var user_id = Math.floor(Math.random() * 1000000000);
+
+
+// initilize
+window.onload = function () {
+    ctxIn.fillStyle = "#87ceeb";
+    ctxIn.fillRect(0, 0, cvsIn.width, 300);
+    ctxIn.fillStyle = "#567d46";
+    ctxIn.fillRect(0, 300, cvsIn.width, 512);
+
+    ctxIn.color = "#b0d49b";
+    ctxIn.lineWidth = 30;
+    ctxIn.lineJoin = ctxIn.lineCap = 'round';
+}
+
+
+// add cavas events
+cvsIn.addEventListener("mousedown", function (e) {
+    if (e.button == 0) {
+        let rect = e.target.getBoundingClientRect();
+        let x = e.clientX - rect.left;
+        let y = e.clientY - rect.top;
+        mouselbtn = true;
+        ctxIn.beginPath();
+        ctxIn.moveTo(x, y);
+    }
+    else if (e.button == 2) {
+        onClear();  // right click for clear input
+    }
+});
+
+cvsIn.addEventListener("mouseup", function (e) {
+    if (e.button == 0) {
+        mouselbtn = false;
+        move_range = domainSlider.value;
+        onRecognition(move_range);
+    }
+});
+cvsIn.addEventListener("mousemove", function (e) {
+    let rect = e.target.getBoundingClientRect();
+    let x = e.clientX - rect.left;
+    let y = e.clientY - rect.top;
+    if (mouselbtn) {
+        ctxIn.lineTo(x, y);
+        ctxIn.strokeStyle = ctxIn.color;
+        ctxIn.stroke();
+        if (((new Date).getTime() - current_time) >= 400) {
+            move_range = domainSlider.value;
+            onRecognition(move_range);
+            current_time = (new Date).getTime();
+        }
+
+    }
+});
+
+cvsIn.addEventListener("touchstart", function (e) {
+    // for touch device
+    if (e.targetTouches.length == 1) {
+        let rect = e.target.getBoundingClientRect();
+        let touch = e.targetTouches[0];
+        let x = touch.clientX - rect.left;
+        let y = touch.clientY - rect.top;
+        ctxIn.beginPath();
+        ctxIn.moveTo(x, y);
+    }
+});
+
+cvsIn.addEventListener("touchmove", function (e) {
+    // for touch device
+    if (e.targetTouches.length == 1) {
+        let rect = e.target.getBoundingClientRect();
+        let touch = e.targetTouches[0];
+        let x = touch.clientX - rect.left;
+        let y = touch.clientY - rect.top;
+        ctxIn.lineTo(x, y);
+        ctxIn.strokeStyle = ctxIn.color;
+        ctxIn.stroke();
+        e.preventDefault();
+    }
+});
+
+cvsIn.addEventListener("touchend", function (e) {
+    // for touch device
+    move_range = domainSlider.value;
+    onRecognition(move_range);
+});
+
+// prevent display the contextmenu
+cvsIn.addEventListener('contextmenu', function (e) {
+    e.preventDefault();
+});
+
+document.getElementById("clearbtn").onclick = onClear;
+function onClear() {
+    mouselbtn = false;
+    ctxIn.clearRect(0, 0, 512, 512);
+    ctxIn.fillStyle = "#87ceeb";
+    ctxIn.fillRect(0, 0, cvsIn.width, 300);
+    ctxIn.fillStyle = "#567d46";
+    ctxIn.fillRect(0, 300, cvsIn.width, 512);
+}
+
+
+document.getElementById("random_pick").addEventListener("click", function () {
+    //ctxIn.color = "#F5F5F5";
+    onRecognition_random();
+});
+
+
+document.getElementById("color1").addEventListener("click", function () {
+    //ctxIn.color = "#D5D5D5";
+    ctxIn.color = "#87ceeb";
+});
+document.getElementById("color2").addEventListener("click", function () {
+    //ctxIn.color = "#696969";
+    ctxIn.color = "#9b7653"
+});
+document.getElementById("color3").addEventListener("click", function () {
+    //ctxIn.color = "#676767";
+    ctxIn.color = "#b0d49b"
+});
+document.getElementById("color4").addEventListener("click", function () {
+    //ctxIn.color = "#F5F5F5";
+    ctxIn.color = "#5abcd8"
+});
+document.getElementById("color5").addEventListener("click", function () {
+    //ctxIn.color = "#F5F5F5";
+    ctxIn.color = "#C1BEBA"
+});
+document.getElementById("color6").addEventListener("click", function () {
+    ctxIn.color = "#5A4D41"
+});
+document.getElementById("color7").addEventListener("click", function () {
+    ctxIn.color = "#567d46"
+});
+document.getElementById("color8").addEventListener("click", function () {
+    ctxIn.color = "#42692f"
+});
+
+document.getElementById("color9").addEventListener("click", function () {
+    ctxIn.color = "#1577be"
+});
+//document.getElementById("color10").addEventListener("click", function(){
+//ctxIn.color = "#676767";
+//    ctxIn.color = "#808080"
+//});
+document.getElementById("color11").addEventListener("click", function () {
+    //ctxIn.color = "#F5F5F5";
+    ctxIn.color = "#3a2e27"
+});
+document.getElementById("color12").addEventListener("click", function () {
+    //ctxIn.color = "#F5F5F5";
+    ctxIn.color = "#4D415A"
+});
+//document.getElementById("color13").addEventListener("click", function(){
+//    ctxIn.color = "#74cc8c"
+//});
+document.getElementById("color14").addEventListener("click", function () {
+    ctxIn.color = "#FDDA16"
+});
+document.getElementById("color15").addEventListener("click", function () {
+    ctxIn.color = "#d0cccc"
+});
+
+var brushSlider = document.getElementById("brushSlider");
+ctxIn.lineWidth = brushSlider.value;
+
+brushSlider.addEventListener("change", function () {
+    ctxIn.lineWidth = brushSlider.value;
+});
+
+
+
+
+var move_range = 3;//domainSlider_1.value;
+
+
+document.getElementById('style').addEventListener('change', function (event) {
+    domainSlider.value = 3;
+    onRecognition(domainSlider.value);
+})
+
+
+domainSlider.addEventListener("change", function () {
+    // style.value = "ink";
+    move_range = domainSlider.value;
+    onRecognition(move_range);
+});
+
+
+// post data to server for recognition
+function onRecognition(range) {
+    console.time("predict");
+
+    $.ajax({
+        url: './predict',
+        type: 'POST',
+        data: JSON.stringify({
+            img: cvsIn.toDataURL("image/png").replace('data:image/png;base64,', ''),
+            model: style.value,
+            move_range: range,
+            user_id: user_id
+        }),
+        contentType: 'application/json',
+    }).done(function (data) {
+        drawImgToCanvas("outputimg", data)
+
+    }).fail(function (XMLHttpRequest, textStatus, errorThrown) {
+        console.log(XMLHttpRequest);
+        alert("error");
+    })
+
+    console.timeEnd("time");
+}
+
+function onRecognition_random(range) {
+    console.time("predict");
+
+    $.ajax({
+        url: './predict_random',
+        type: 'POST',
+        data: JSON.stringify({
+            img: cvsIn.toDataURL("image/png").replace('data:image/png;base64,', ''),
+            model: style.value,
+            move_range: range,
+            user_id: user_id
+        }),
+        contentType: 'application/json',
+    }).done(function (data) {
+        drawImgToCanvas("outputimg", data)
+
+    }).fail(function (XMLHttpRequest, textStatus, errorThrown) {
+        console.log(XMLHttpRequest);
+        alert("error");
+    })
+
+    console.timeEnd("time");
+}
+
+function drawImgToCanvas(canvasId, b64Img) {
+    let canvas = document.getElementById(canvasId);
+    let ctx = canvas.getContext('2d');
+    let img = new Image();
+    img.src = "data:image/png;base64," + b64Img;
+    img.onload = function () {
+        ctx.drawImage(img, 0, 0, img.width, img.height, 0, 0, canvas.width, canvas.height);
+    }
+}

static/style.css

@@ -0,0 +1,36 @@
+.common {
+  text-align: center;
+}
+
+.boxitem1 {
+  display: inline-block;
+  vertical-align: left;
+}
+
+select {
+  font-size: 1.7em;
+  border: 1px;
+}
+
+button {
+  border-radius: 10px;
+}
+
+.boxitem2 {
+  display: inline-block;
+  vertical-align: right;
+}
+.boxitem {
+  display: inline-block;
+  vertical-align: top;
+}
+
+#inputimg {
+  vertical-align: left;
+  border: solid 1px black;
+}
+
+#outputimg {
+  vertical-align: left;
+  border: solid 1px black;
+}

templates/index.html

@@ -0,0 +1,124 @@
+<!DOCTYPE html>
+<html>
+
+<head>
+  <link rel="stylesheet" href="https://cdn.rawgit.com/Chalarangelo/mini.css/v2.3.7/dist/mini-default.min.css">
+  <link rel="stylesheet" href="./static/style.css">
+  <meta charset="UTF-8">
+  <title>Label to Art Demo_V_0.7 </title>
+</head>
+
+<div class="common">
+
+  <body>
+    <div class="row">
+      <div class="col-sm-12 col-md-10 col-md-offset-1">
+        &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
+        <div class="boxitem1">
+          <label for='brushSlider' style='font-size: 1.7em; font-weight: bold'>Stroke Width </label>
+          <input type="range" name="brushsize" min="0" max="100" id="brushSlider" step="1" value="30"
+            onchange="this.setAttribute('value',this.value);">
+        </div>
+
+        &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
+        &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
+        &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
+        &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
+        <div class="boxitem2">
+          &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
+          <label for="style" style='font-size: 1.7em; font-weight: bold'>Domain:</label>
+          <select id="style">
+            <option style='font-size: 1.7em; font-weight: bold' value="ink">Ink Wash</option>
+            <option style='font-size: 1.7em; font-weight: bold' value="monet">Monet</option>
+            <option style='font-size: 1.7em; font-weight: bold' value="vangogh">Van Gogh</option>
+            <option style='font-size: 1.7em; font-weight: bold' value="water">WaterColor</option>
+          </select>
+        </div>
+
+        <div class="row">
+          <div class="col-sm-12 col-md-10 col-md-offset-1">
+            <div class="boxitem1">
+              <label for='domainSlider' style='font-size: 1.7em; font-weight: bold'> Style Strength </label>
+              <input type="range" name="style range" min="1" max="5" id="domainSlider" step="0.2" value="3"
+                onchange="this.setAttribute('value',this.value);">
+            </div>
+            &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
+            &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
+            &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
+            &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
+            &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
+            &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
+            &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
+            &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;
+            <div class="boxitem2">
+              <button id="random_pick" class="primary"
+                style="background-color:#003449; margin-right:0px; font-size: 1.9em; font-weight: bold; border-radius:15px;">
+                Random </button>
+            </div>
+          </div>
+          <br>
+
+        </div>
+        <div class="divider"></div>
+        <div class="boxitem1">
+          <canvas id="inputimg" width="512" height="512" style="border:5px solid #ffffff;"></canvas>
+        </div>
+        <div class="boxitem2">
+          <canvas id="outputimg" width="512" height="512" style="border:5px solid #ffffff;"></canvas>
+        </div>
+      </div>
+    </div>
+
+
+    <div class="boxitem">
+      <button id="color1" class="primary"
+        style="background-color:#87ceeb; margin-right:0px; font-size: 1.9em; font-weight: bold; border-radius:15px;">
+        Sky </button>
+      <button id="color2" class="primary"
+        style="background-color:#9b7653; margin-right:0px; font-size: 1.9em; font-weight: bold; border-radius:15px;">
+        Dirt </button>
+      <button id="color3" class="primary"
+        style="background-color:#b0d49b; margin-right:0px; font-size: 1.9em; font-weight: bold; border-radius:15px;">
+        Mountain </button>
+      <button id="color4" class="primary"
+        style="background-color:#5abcd8; margin-right:0px; font-size: 1.9em; font-weight: bold; border-radius:15px;">
+        River </button>
+      <button id="color5" class="primary"
+        style="background-color:#C1BEBA; margin-right:0px; font-size: 1.9em; font-weight: bold; border-radius:15px;">
+        Clouds </button>
+      <button id="color6" class="primary"
+        style="background-color:#5A4D41; margin-right:0px; font-size: 1.9em; font-weight: bold; border-radius:15px;">
+        Rock </button>
+      <button id="color7" class="primary"
+        style="background-color:#567d46; margin-right:0px; font-size: 1.9em; font-weight: bold; border-radius:15px;">
+        Grass </button>
+      <button id="color8" class="primary"
+        style="background-color:#42692f; margin-right:0px; font-size: 1.9em; font-weight: bold; border-radius:15px;">
+        Tree </button>
+    </div>
+    <br>
+    <div class="boxitem">
+      <button id="color9" class="primary"
+        style="background-color:#1577be; margin-right:0px; font-size: 1.9em; font-weight: bold; border-radius:15px;">
+        Sea </button>
+      <button id="color11" class="primary"
+        style="background-color:#3a2e27; margin-right:0px; font-size: 1.9em; font-weight: bold; border-radius:15px;">
+        Ground </button>
+      <button id="color12" class="primary"
+        style="background-color:#4D415A; margin-right:0px; font-size: 1.9em; font-weight: bold; border-radius:15px;">
+        Hill </button>
+      <button id="color14" class="primary"
+        style="background-color:#FDDA16; margin-right:0px; font-size: 1.9em; font-weight: bold; border-radius:15px;">
+        Road </button>
+      <button id="color15" class="primary"
+        style="background-color:#d0cccc; margin-right:0px; font-size: 1.9em; font-weight: bold; border-radius:15px;">
+        Snow </button>
+      <button id="clearbtn" class="primary"
+        style="margin-right:0px; font-size: 1.9em; font-weight: bold; border-radius:15px;"> Clear </button>
+
+    </div>
+
+    <script src="./static/index.js"></script>
+    <script src="//ajax.googleapis.com/ajax/libs/jquery/3.3.1/jquery.min.js"></script>
+    <script src="//d3js.org/d3.v5.min.js"></script>
+  </body>

utils/boundaries_amp_52/artwork_ink_boundary/boundary.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63f02c41f68b7fcd603e158fbd4542ebbe3b42d5b25b64a5e1e1e0699e1286c0
+size 1152

utils/boundaries_amp_52/artwork_ink_boundary/log.txt

@@ -0,0 +1,12 @@
+[2021-09-15 00:57:58,220][INFO] Loading latent codes.
+[2021-09-15 00:57:58,233][INFO] Loading attribute scores.
+[2021-09-15 00:57:58,234][INFO] Filtering training data.
+[2021-09-15 00:57:58,234][INFO] Sorting scores to get positive and negative samples.
+[2021-09-15 00:57:58,246][INFO] Spliting training and validation sets:
+[2021-09-15 00:57:58,337][INFO]   Training: 4200 positive, 4200 negative.
+[2021-09-15 00:57:58,342][INFO]   Validation: 1800 positive, 1800 negative.
+[2021-09-15 00:57:58,353][INFO]   Remaining: 4255 positive, 23745 negative.
+[2021-09-15 00:57:58,356][INFO] Training boundary.
+[2021-09-15 00:57:59,433][INFO] Finish training.
+[2021-09-15 00:57:59,711][INFO] Accuracy for validation set: 3596 / 3600 = 0.998889
+[2021-09-15 00:58:01,331][INFO] Accuracy for remaining set: 27037 / 28000 = 0.965607

utils/boundaries_amp_52/artwork_monet_boundary/boundary.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:477c0e8af08ccb05d8ddb79e38b5c71cbc856bb4ebe0adf80194de48cf1510ad
+size 1152

utils/boundaries_amp_52/artwork_monet_boundary/log.txt

@@ -0,0 +1,12 @@
+[2021-09-15 00:57:58,250][INFO] Loading latent codes.
+[2021-09-15 00:57:58,265][INFO] Loading attribute scores.
+[2021-09-15 00:57:58,266][INFO] Filtering training data.
+[2021-09-15 00:57:58,266][INFO] Sorting scores to get positive and negative samples.
+[2021-09-15 00:57:58,281][INFO] Spliting training and validation sets:
+[2021-09-15 00:57:58,393][INFO]   Training: 4200 positive, 4200 negative.
+[2021-09-15 00:57:58,398][INFO]   Validation: 1800 positive, 1800 negative.
+[2021-09-15 00:57:58,400][INFO]   Remaining: 5854 positive, 22146 negative.
+[2021-09-15 00:57:58,407][INFO] Training boundary.
+[2021-09-15 00:57:59,699][INFO] Finish training.
+[2021-09-15 00:57:59,912][INFO] Accuracy for validation set: 3549 / 3600 = 0.985833
+[2021-09-15 00:58:01,556][INFO] Accuracy for remaining set: 23849 / 28000 = 0.851750

utils/boundaries_amp_52/artwork_vangogh_boundary/boundary.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b616bf764d3ea1c85b9f852b5320b90fcb6ff5d90f80d22ef8a7e5c3a47a5959
+size 1152

utils/boundaries_amp_52/artwork_vangogh_boundary/log.txt

@@ -0,0 +1,12 @@
+[2021-09-15 00:57:58,217][INFO] Loading latent codes.
+[2021-09-15 00:57:58,230][INFO] Loading attribute scores.
+[2021-09-15 00:57:58,231][INFO] Filtering training data.
+[2021-09-15 00:57:58,231][INFO] Sorting scores to get positive and negative samples.
+[2021-09-15 00:57:58,246][INFO] Spliting training and validation sets:
+[2021-09-15 00:57:58,279][INFO]   Training: 4200 positive, 4200 negative.
+[2021-09-15 00:57:58,281][INFO]   Validation: 1800 positive, 1800 negative.
+[2021-09-15 00:57:58,281][INFO]   Remaining: 3401 positive, 24599 negative.
+[2021-09-15 00:57:58,281][INFO] Training boundary.
+[2021-09-15 00:57:59,347][INFO] Finish training.
+[2021-09-15 00:57:59,551][INFO] Accuracy for validation set: 3596 / 3600 = 0.998889
+[2021-09-15 00:58:01,098][INFO] Accuracy for remaining set: 23785 / 28000 = 0.849464

utils/boundaries_amp_52/artwork_water_boundary/boundary.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f0827fe3ea8f6361bc298c968623c2ff2e33e5183f4d9df18c118f751c10c713
+size 1152

utils/boundaries_amp_52/artwork_water_boundary/log.txt

@@ -0,0 +1,12 @@
+[2021-09-15 00:57:58,226][INFO] Loading latent codes.
+[2021-09-15 00:57:58,244][INFO] Loading attribute scores.
+[2021-09-15 00:57:58,245][INFO] Filtering training data.
+[2021-09-15 00:57:58,245][INFO] Sorting scores to get positive and negative samples.
+[2021-09-15 00:57:58,263][INFO] Spliting training and validation sets:
+[2021-09-15 00:57:58,390][INFO]   Training: 4200 positive, 4200 negative.
+[2021-09-15 00:57:58,393][INFO]   Validation: 1800 positive, 1800 negative.
+[2021-09-15 00:57:58,398][INFO]   Remaining: 4465 positive, 23535 negative.
+[2021-09-15 00:57:58,401][INFO] Training boundary.
+[2021-09-15 00:57:59,812][INFO] Finish training.
+[2021-09-15 00:58:00,021][INFO] Accuracy for validation set: 3584 / 3600 = 0.995556
+[2021-09-15 00:58:01,830][INFO] Accuracy for remaining set: 24271 / 28000 = 0.866821

utils/umap_utils.py

@@ -0,0 +1,99 @@
+import numpy as np
+import torch
+import copy
+import os
+import numpy as np
+from sklearn import svm
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+
+def linear_interpolate(latent_code, boundary, start_distance=-3, end_distance=3, steps=10):
+    """Manipulates the given latent code with respect to a particular boundary.
+
+    Basically, this function takes a latent code and a boundary as inputs, and
+    outputs a collection of manipulated latent codes. For example, let `steps` to
+    be 10, then the input `latent_code` is with shape [1, latent_space_dim], input
+    `boundary` is with shape [1, latent_space_dim] and unit norm, the output is
+    with shape [10, latent_space_dim]. The first output latent code is
+    `start_distance` away from the given `boundary`, while the last output latent
+    code is `end_distance` away from the given `boundary`. Remaining latent codes
+    are linearly interpolated.
+
+    Input `latent_code` can also be with shape [1, num_layers, latent_space_dim]
+    to support W+ space in Style GAN. In this case, all features in W+ space will
+    be manipulated same as each other. Accordingly, the output will be with shape
+    [10, num_layers, latent_space_dim].
+
+    NOTE: Distance is sign sensitive.
+
+    Args:
+        latent_code: The input latent code for manipulation.
+        boundary: The semantic boundary as reference.
+        start_distance: The distance to the boundary where the manipulation starts.
+        (default: -3.0)
+        end_distance: The distance to the boundary where the manipulation ends.
+        (default: 3.0)
+        steps: Number of steps to move the latent code from start position to end
+        position. (default: 10)
+    """
+    assert latent_code.shape[0] == 1 and boundary.shape[0] == 1 and len(boundary.shape) == 2 and boundary.shape[1] == latent_code.shape[-1]
+
+    linspace = np.linspace(start_distance, end_distance, steps)
+    if len(latent_code.shape) == 2:
+        linspace = linspace - latent_code.dot(boundary.T)
+        linspace = linspace.reshape(-1, 1).astype(np.float32)
+        return latent_code + linspace * boundary
+    if len(latent_code.shape) == 3:
+        linspace = linspace.reshape(-1, 1, 1).astype(np.float32)
+        return latent_code + linspace * boundary.reshape(1, 1, -1)
+    raise ValueError(
+        f"Input `latent_code` should be with shape "
+        f"[1, latent_space_dim] or [1, N, latent_space_dim] for "
+        f"W+ space in Style GAN!\n"
+        f"But {latent_code.shape} is received."
+    )
+
+
+def get_code(domain, boundaries):
+    if domain == "ink":
+        domain = 0
+    elif domain == "monet":
+        domain = 1
+    elif domain == "vangogh":
+        domain = 2
+    elif domain == "water":
+        domain = 3
+
+    res = np.array(torch.randn(1, 256, dtype=torch.float32))
+    # res = linear_interpolate(res, boundaries[domain], end_distance=3, steps=3)[-1:]
+    res = torch.Tensor(res).cuda() if torch.cuda.is_available() else torch.Tensor(res)
+    return res
+
+
+def modify_code(code, boundaries, domain, range):
+    if domain == "ink":
+        domain = 0
+    elif domain == "monet":
+        domain = 1
+    elif domain == "vangogh":
+        domain = 2
+    elif domain == "water":
+        domain = 3
+    # print(domain, range)
+    if range == 0:
+        return code
+    else:
+        res = np.array(code.cpu().detach().numpy())
+        res = linear_interpolate(res, boundaries[domain], end_distance=range, steps=3)[-1:]
+        res = torch.Tensor(res).cuda() if torch.cuda.is_available() else torch.Tensor(res)
+        return res
+
+
+def load_boundries():
+    domains = ["ink", "monet", "vangogh", "water"]
+    domains.sort()
+    boundaries = [
+        np.load(os.path.join(os.path.dirname(__file__), "boundaries_amp_52/artwork_" + domain + "_boundary/boundary.npy")) for domain in domains
+    ]
+    return boundaries