first commit

85 files changed, 1982 insertions, 5 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100755
index 0000000..681efd0
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,40 @@
+debug*
+checkpoints/
+results/
+build/
+dist/
+torch.egg-info/
+*/**/__pycache__
+torch/version.py
+torch/csrc/generic/TensorMethods.cpp
+torch/lib/*.so*
+torch/lib/*.dylib*
+torch/lib/*.h
+torch/lib/build
+torch/lib/tmp_install
+torch/lib/include
+torch/lib/torch_shm_manager
+torch/csrc/cudnn/cuDNN.cpp
+torch/csrc/nn/THNN.cwrap
+torch/csrc/nn/THNN.cpp
+torch/csrc/nn/THCUNN.cwrap
+torch/csrc/nn/THCUNN.cpp
+torch/csrc/nn/THNN_generic.cwrap
+torch/csrc/nn/THNN_generic.cpp
+torch/csrc/nn/THNN_generic.h
+docs/src/**/*
+test/data/legacy_modules.t7
+test/data/gpu_tensors.pt
+test/htmlcov
+test/.coverage
+*/*.pyc
+*/**/*.pyc
+*/**/**/*.pyc
+*/**/**/**/*.pyc
+*/**/**/**/**/*.pyc
+*/*.so*
+*/**/*.so*
+*/**/*.dylib*
+test/data/legacy_serialized.pt
+*.DS_Store
+*~
diff --git a/LICENSE.txt b/LICENSE.txt
new file mode 100755
index 0000000..7406cd7
--- /dev/null
+++ b/LICENSE.txt
@@ -0,0 +1,43 @@
+Copyright (C) 2017 NVIDIA Corporation. Ting-Chun Wang, Ming-Yu Liu, Jun-Yan Zhu.
+All rights reserved. 
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+
+Permission to use, copy, modify, and distribute this software and its documentation 
+for any non-commercial purpose is hereby granted without fee, provided that the above 
+copyright notice appear in all copies and that both that copyright notice and this 
+permission notice appear in supporting documentation, and that the name of the author 
+not be used in advertising or publicity pertaining to distribution of the software 
+without specific, written prior permission.
+
+THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE, INCLUDING ALL 
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR ANY PARTICULAR PURPOSE. 
+IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL 
+DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, 
+WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING 
+OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+
+--------------------------- LICENSE FOR pytorch-CycleGAN-and-pix2pix ----------------
+Copyright (c) 2017, Jun-Yan Zhu and Taesung Park
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
diff --git a/README.md b/README.md
index fc13ade..c560c46 100755
--- a/README.md
+++ b/README.md
@@ -10,10 +10,6 @@ Pytorch implementation of our method for high-resolution (e.g. 2048x1024) photor
  <sup>1</sup>NVIDIA Corporation, <sup>2</sup>UC Berkeley  

  In arxiv, 2017.  

 

-## Release notice

-The code is ready to publish but still under final approval process. It should be approved in a couple of days.<br>

-If you want to get notified once the code is released, please subscribe [here](https://tcwang0509.github.io/pix2pixHD/subscribe.html).

-

 ## Image-to-image translation at 2k/1k resolution

 - Our label-to-streetview results

 <p align='center'>  

@@ -53,7 +49,69 @@ If you want to get notified once the code is released, please subscribe [here](h
   <img src='imgs/face_short.gif' width='490'/>

 </p>

 

-### Citation

+## Prerequisites

+- Linux or macOS

+- Python 2 or 3

+- NVIDIA GPU (12G or 24G memory) + CUDA cuDNN

+

+## Getting Started

+### Installation

+- Install PyTorch and dependencies from http://pytorch.org

+- Install python libraries [dominate](https://github.com/Knio/dominate).

+```bash

+pip install dominate

+```

+- Clone this repo:

+```bash

+git clone https://github.com/NVIDIA/pix2pixHD

+cd pix2pixHD

+```

+

+

+### Testing

+- A few example Cityscapes test images are included in the `datasets` folder.

+- Please download the pre-trained Cityscapes model from [here](https://drive.google.com/file/d/1h9SykUnuZul7J3Nbms2QGH1wa85nbN2-/view?usp=sharing) (google drive link), and put it under `./checkpoints/label2city_1024p/`

+- Test the model (`bash ./scripts/test_1024p.sh`):

+```bash

+#!./scripts/test_1024p.sh

+python test.py --name label2city_1024p --netG local --ngf 32 --resize_or_crop none

+```

+The test results will be saved to a html file here: `./results/label2city_1024p/test_latest/index.html`.

+

+More example scripts can be found in the `scripts` directory.

+

+

+### Dataset

+- We use the Cityscapes dataset. To train a model on the full dataset, please download it from the [official website](https://www.cityscapes-dataset.com/) (registration required).

+After downloading, please put it under the `datasets` folder in the same way the example images are provided.

+

+

+### Training

+- Train a model at 1024 x 512 resolution (`bash ./scripts/train_512p.sh`):

+```bash

+#!./scripts/train_512p.sh

+python train.py --name label2city_512p

+```

+- To view training results, please checkout intermediate results in `./checkpoints/label2city_512p/web/index.html`.

+If you have tensorflow installed, you can see tensorboard logs in `./checkpoints/label2city_512p/logs` by adding `--tf_log` to the training scripts.

+

+### Multi-GPU training

+- Train a model using multiple GPUs (`bash ./scripts/train_512p_multigpu.sh`):

+```bash

+#!./scripts/train_512p_multigpu.sh

+python train.py --name label2city_512p --batchSize 8 --gpu_ids 0,1,2,3,4,5,6,7

+```

+Note: this is not tested and we trained our model using single GPU only. Please use at your own discretion.

+

+### Training at full resolution

+- To train the images at full resolution (2048 x 1024) requires a GPU with 24G memory (`bash ./scripts/train_1024p_24G.sh`).

+If only GPUs with 12G memory are available, please use the 12G script (`bash ./scripts/train_1024p_12G.sh`), which will crop the images during training. Performance is not guaranteed using this script.

+

+## More Training/test Details

+- Flags: see `options/train_options.py` and `options/base_options.py` for all the training flags; see `options/test_options.py` and `options/base_options.py` for all the test flags.

+

+

+## Citation

 

 If you find this useful for your research, please use the following.

 

@@ -65,3 +123,6 @@ If you find this useful for your research, please use the following.
   year={2017}

 }

 ```

+

+## Acknowledgments

+This code borrows heavily from [pytorch-CycleGAN-and-pix2pix](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix).

diff --git a/_config.yml b/_config.yml
new file mode 100755
index 0000000..2f7efbe
--- /dev/null
+++ b/_config.yml
@@ -0,0 +1 @@
+theme: jekyll-theme-minimal
\ No newline at end of file
diff --git a/data/__init__.py b/data/__init__.py
new file mode 100755
index 0000000..e69de29
--- /dev/null
+++ b/data/__init__.py
diff --git a/data/aligned_dataset.py b/data/aligned_dataset.py
new file mode 100755
index 0000000..50390f3
--- /dev/null
+++ b/data/aligned_dataset.py
@@ -0,0 +1,76 @@
+### Copyright (C) 2017 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.path
+import random
+import torchvision.transforms as transforms
+import torch
+from data.base_dataset import BaseDataset, get_params, get_transform, normalize
+from data.image_folder import make_dataset
+from PIL import Image
+import numpy as np
+
+class AlignedDataset(BaseDataset):
+    def initialize(self, opt):
+        self.opt = opt
+        self.root = opt.dataroot    
+
+        ### label maps         
+        self.dir_label = os.path.join(opt.dataroot, opt.phase + '_label')              
+        self.label_paths = sorted(make_dataset(self.dir_label))
+
+        ### real images
+        if opt.isTrain:
+            self.dir_image = os.path.join(opt.dataroot, opt.phase + '_img')  
+            self.image_paths = sorted(make_dataset(self.dir_image))
+
+        ### instance maps
+        if not opt.no_instance:
+            self.dir_inst = os.path.join(opt.dataroot, opt.phase + '_inst')
+            self.inst_paths = sorted(make_dataset(self.dir_inst))
+
+        ### load precomputed instance-wise encoded features
+        if opt.load_features:                              
+            self.dir_feat = os.path.join(opt.dataroot, opt.phase + '_feat')
+            print('----------- loading features from %s ----------' % self.dir_feat)
+            self.feat_paths = sorted(make_dataset(self.dir_feat))
+
+        self.dataset_size = len(self.label_paths) 
+      
+    def __getitem__(self, index):        
+        ### label maps
+        label_path = self.label_paths[index]              
+        label = Image.open(label_path)        
+        params = get_params(self.opt, label.size)          
+        transform_label = get_transform(self.opt, params, method=Image.NEAREST, normalize=False)
+        label_tensor = transform_label(label) * 255.0
+
+        image_tensor = inst_tensor = feat_tensor = 0
+        ### real images
+        if self.opt.isTrain:
+            image_path = self.image_paths[index]   
+            image = Image.open(image_path).convert('RGB')
+            transform_image = get_transform(self.opt, params)      
+            image_tensor = transform_image(image)
+
+        ### if using instance maps        
+        if not self.opt.no_instance:
+            inst_path = self.inst_paths[index]
+            inst = Image.open(inst_path)
+            inst_tensor = transform_label(inst)
+
+            if self.opt.load_features:
+                feat_path = self.feat_paths[index]            
+                feat = Image.open(feat_path).convert('RGB')
+                norm = normalize()
+                feat_tensor = norm(transform_label(feat))                            
+
+        input_dict = {'label': label_tensor, 'inst': inst_tensor, 'image': image_tensor, 
+                      'feat': feat_tensor, 'path': label_path}
+
+        return input_dict
+
+    def __len__(self):
+        return len(self.label_paths)
+
+    def name(self):
+        return 'AlignedDataset'
\ No newline at end of file
diff --git a/data/base_data_loader.py b/data/base_data_loader.py
new file mode 100755
index 0000000..0e1deb5
--- /dev/null
+++ b/data/base_data_loader.py
@@ -0,0 +1,14 @@
+
+class BaseDataLoader():
+    def __init__(self):
+        pass
+    
+    def initialize(self, opt):
+        self.opt = opt
+        pass
+
+    def load_data():
+        return None
+
+        
+        
diff --git a/data/base_dataset.py b/data/base_dataset.py
new file mode 100755
index 0000000..038d3d2
--- /dev/null
+++ b/data/base_dataset.py
@@ -0,0 +1,92 @@
+### Copyright (C) 2017 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.utils.data as data
+from PIL import Image
+import torchvision.transforms as transforms
+import numpy as np
+import random
+
+class BaseDataset(data.Dataset):
+    def __init__(self):
+        super(BaseDataset, self).__init__()
+
+    def name(self):
+        return 'BaseDataset'
+
+    def initialize(self, opt):
+        pass
+
+def get_params(opt, size):
+    w, h = size
+    new_h = h
+    new_w = w
+    if opt.resize_or_crop == 'resize_and_crop':
+        new_h = new_w = opt.loadSize            
+    elif opt.resize_or_crop == 'scale_width_and_crop':
+        new_w = opt.loadSize
+        new_h = opt.loadSize * h // w
+
+    x = random.randint(0, np.maximum(0, new_w - opt.fineSize))
+    y = random.randint(0, np.maximum(0, new_h - opt.fineSize))
+    
+    flip = random.random() > 0.5
+    return {'crop_pos': (x, y), 'flip': flip}
+
+def get_transform(opt, params, method=Image.BICUBIC, normalize=True):
+    transform_list = []
+    if 'resize' in opt.resize_or_crop:
+        osize = [opt.loadSize, opt.loadSize]
+        transform_list.append(transforms.Scale(osize, method))   
+    elif 'scale_width' in opt.resize_or_crop:
+        transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.loadSize, method)))
+        
+    if 'crop' in opt.resize_or_crop:
+        transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], opt.fineSize)))
+
+    if opt.resize_or_crop == 'none':
+        base = float(2 ** opt.n_downsample_global)
+        if opt.netG == 'local':
+            base *= (2 ** opt.n_local_enhancers)
+        transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base, method)))
+
+    if opt.isTrain and not opt.no_flip:
+        transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip'])))
+
+    transform_list += [transforms.ToTensor()]
+
+    if normalize:
+        transform_list += [transforms.Normalize((0.5, 0.5, 0.5),
+                                                (0.5, 0.5, 0.5))]
+    return transforms.Compose(transform_list)
+
+def normalize():    
+    return transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+
+def __make_power_2(img, base, method=Image.BICUBIC):
+    ow, oh = img.size        
+    h = int(round(oh / base) * base)
+    w = int(round(ow / base) * base)
+    if (h == oh) and (w == ow):
+        return img
+    return img.resize((w, h), method)
+
+def __scale_width(img, target_width, method=Image.BICUBIC):
+    ow, oh = img.size
+    if (ow == target_width):
+        return img    
+    w = target_width
+    h = int(target_width * oh / ow)    
+    return img.resize((w, h), method)
+
+def __crop(img, pos, size):
+    ow, oh = img.size
+    x1, y1 = pos
+    tw = th = size
+    if (ow > tw or oh > th):        
+        return img.crop((x1, y1, x1 + tw, y1 + th))
+    return img
+
+def __flip(img, flip):
+    if flip:
+        return img.transpose(Image.FLIP_LEFT_RIGHT)
+    return img
diff --git a/data/custom_dataset_data_loader.py b/data/custom_dataset_data_loader.py
new file mode 100755
index 0000000..0b98254
--- /dev/null
+++ b/data/custom_dataset_data_loader.py
@@ -0,0 +1,31 @@
+import torch.utils.data
+from data.base_data_loader import BaseDataLoader
+
+
+def CreateDataset(opt):
+    dataset = None
+    from data.aligned_dataset import AlignedDataset
+    dataset = AlignedDataset()
+
+    print("dataset [%s] was created" % (dataset.name()))
+    dataset.initialize(opt)
+    return dataset
+
+class CustomDatasetDataLoader(BaseDataLoader):
+    def name(self):
+        return 'CustomDatasetDataLoader'
+
+    def initialize(self, opt):
+        BaseDataLoader.initialize(self, opt)
+        self.dataset = CreateDataset(opt)
+        self.dataloader = torch.utils.data.DataLoader(
+            self.dataset,
+            batch_size=opt.batchSize,
+            shuffle=not opt.serial_batches,
+            num_workers=int(opt.nThreads))
+
+    def load_data(self):
+        return self.dataloader
+
+    def __len__(self):
+        return min(len(self.dataset), self.opt.max_dataset_size)
diff --git a/data/data_loader.py b/data/data_loader.py
new file mode 100755
index 0000000..2a4433a
--- /dev/null
+++ b/data/data_loader.py
@@ -0,0 +1,7 @@
+
+def CreateDataLoader(opt):
+    from data.custom_dataset_data_loader import CustomDatasetDataLoader
+    data_loader = CustomDatasetDataLoader()
+    print(data_loader.name())
+    data_loader.initialize(opt)
+    return data_loader
diff --git a/data/image_folder.py b/data/image_folder.py
new file mode 100755
index 0000000..16a447c
--- /dev/null
+++ b/data/image_folder.py
@@ -0,0 +1,68 @@
+###############################################################################
+# Code from
+# https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py
+# Modified the original code so that it also loads images from the current
+# directory as well as the subdirectories
+###############################################################################
+
+import torch.utils.data as data
+
+from PIL import Image
+import os
+import os.path
+
+IMG_EXTENSIONS = [
+    '.jpg', '.JPG', '.jpeg', '.JPEG',
+    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tiff'
+]
+
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def make_dataset(dir):
+    images = []
+    assert os.path.isdir(dir), '%s is not a valid directory' % dir
+
+    for root, _, fnames in sorted(os.walk(dir)):
+        for fname in fnames:
+            if is_image_file(fname):
+                path = os.path.join(root, fname)
+                images.append(path)
+
+    return images
+
+
+def default_loader(path):
+    return Image.open(path).convert('RGB')
+
+
+class ImageFolder(data.Dataset):
+
+    def __init__(self, root, transform=None, return_paths=False,
+                 loader=default_loader):
+        imgs = make_dataset(root)
+        if len(imgs) == 0:
+            raise(RuntimeError("Found 0 images in: " + root + "\n"
+                               "Supported image extensions are: " +
+                               ",".join(IMG_EXTENSIONS)))
+
+        self.root = root
+        self.imgs = imgs
+        self.transform = transform
+        self.return_paths = return_paths
+        self.loader = loader
+
+    def __getitem__(self, index):
+        path = self.imgs[index]
+        img = self.loader(path)
+        if self.transform is not None:
+            img = self.transform(img)
+        if self.return_paths:
+            return img, path
+        else:
+            return img
+
+    def __len__(self):
+        return len(self.imgs)
diff --git a/datasets/cityscapes/test_inst/frankfurt_000000_000576_gtFine_instanceIds.png b/datasets/cityscapes/test_inst/frankfurt_000000_000576_gtFine_instanceIds.png
new file mode 100755
index 0000000..01da7ed
--- /dev/null
+++ b/datasets/cityscapes/test_inst/frankfurt_000000_000576_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/test_inst/frankfurt_000000_001236_gtFine_instanceIds.png b/datasets/cityscapes/test_inst/frankfurt_000000_001236_gtFine_instanceIds.png
new file mode 100755
index 0000000..75506bc
--- /dev/null
+++ b/datasets/cityscapes/test_inst/frankfurt_000000_001236_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/test_inst/frankfurt_000000_003357_gtFine_instanceIds.png b/datasets/cityscapes/test_inst/frankfurt_000000_003357_gtFine_instanceIds.png
new file mode 100755
index 0000000..9bd27b0
--- /dev/null
+++ b/datasets/cityscapes/test_inst/frankfurt_000000_003357_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/test_inst/frankfurt_000000_011810_gtFine_instanceIds.png b/datasets/cityscapes/test_inst/frankfurt_000000_011810_gtFine_instanceIds.png
new file mode 100755
index 0000000..df84eee
--- /dev/null
+++ b/datasets/cityscapes/test_inst/frankfurt_000000_011810_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/test_inst/frankfurt_000000_012868_gtFine_instanceIds.png b/datasets/cityscapes/test_inst/frankfurt_000000_012868_gtFine_instanceIds.png
new file mode 100755
index 0000000..ba1f7aa
--- /dev/null
+++ b/datasets/cityscapes/test_inst/frankfurt_000000_012868_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/test_inst/frankfurt_000001_013710_gtFine_instanceIds.png b/datasets/cityscapes/test_inst/frankfurt_000001_013710_gtFine_instanceIds.png
new file mode 100755
index 0000000..d05b7db
--- /dev/null
+++ b/datasets/cityscapes/test_inst/frankfurt_000001_013710_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/test_inst/frankfurt_000001_015328_gtFine_instanceIds.png b/datasets/cityscapes/test_inst/frankfurt_000001_015328_gtFine_instanceIds.png
new file mode 100755
index 0000000..32d62a3
--- /dev/null
+++ b/datasets/cityscapes/test_inst/frankfurt_000001_015328_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/test_inst/frankfurt_000001_023769_gtFine_instanceIds.png b/datasets/cityscapes/test_inst/frankfurt_000001_023769_gtFine_instanceIds.png
new file mode 100755
index 0000000..9eef682
--- /dev/null
+++ b/datasets/cityscapes/test_inst/frankfurt_000001_023769_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/test_inst/frankfurt_000001_028335_gtFine_instanceIds.png b/datasets/cityscapes/test_inst/frankfurt_000001_028335_gtFine_instanceIds.png
new file mode 100755
index 0000000..b1909d5
--- /dev/null
+++ b/datasets/cityscapes/test_inst/frankfurt_000001_028335_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/test_inst/frankfurt_000001_032711_gtFine_instanceIds.png b/datasets/cityscapes/test_inst/frankfurt_000001_032711_gtFine_instanceIds.png
new file mode 100755
index 0000000..ac2e293
--- /dev/null
+++ b/datasets/cityscapes/test_inst/frankfurt_000001_032711_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/test_inst/frankfurt_000001_033655_gtFine_instanceIds.png b/datasets/cityscapes/test_inst/frankfurt_000001_033655_gtFine_instanceIds.png
new file mode 100755
index 0000000..de7328e
--- /dev/null
+++ b/datasets/cityscapes/test_inst/frankfurt_000001_033655_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/test_inst/frankfurt_000001_042733_gtFine_instanceIds.png b/datasets/cityscapes/test_inst/frankfurt_000001_042733_gtFine_instanceIds.png
new file mode 100755
index 0000000..a98d096
--- /dev/null
+++ b/datasets/cityscapes/test_inst/frankfurt_000001_042733_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/test_inst/frankfurt_000001_047552_gtFine_instanceIds.png b/datasets/cityscapes/test_inst/frankfurt_000001_047552_gtFine_instanceIds.png
new file mode 100755
index 0000000..ab569e3
--- /dev/null
+++ b/datasets/cityscapes/test_inst/frankfurt_000001_047552_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/test_inst/frankfurt_000001_054640_gtFine_instanceIds.png b/datasets/cityscapes/test_inst/frankfurt_000001_054640_gtFine_instanceIds.png
new file mode 100755
index 0000000..5f246a6
--- /dev/null
+++ b/datasets/cityscapes/test_inst/frankfurt_000001_054640_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/test_inst/frankfurt_000001_055387_gtFine_instanceIds.png b/datasets/cityscapes/test_inst/frankfurt_000001_055387_gtFine_instanceIds.png
new file mode 100755
index 0000000..2e7d01f
--- /dev/null
+++ b/datasets/cityscapes/test_inst/frankfurt_000001_055387_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/test_label/frankfurt_000000_000576_gtFine_labelIds.png b/datasets/cityscapes/test_label/frankfurt_000000_000576_gtFine_labelIds.png
new file mode 100755
index 0000000..8c9464c
--- /dev/null
+++ b/datasets/cityscapes/test_label/frankfurt_000000_000576_gtFine_labelIds.png
diff --git a/datasets/cityscapes/test_label/frankfurt_000000_001236_gtFine_labelIds.png b/datasets/cityscapes/test_label/frankfurt_000000_001236_gtFine_labelIds.png
new file mode 100755
index 0000000..9f0ca9f
--- /dev/null
+++ b/datasets/cityscapes/test_label/frankfurt_000000_001236_gtFine_labelIds.png
diff --git a/datasets/cityscapes/test_label/frankfurt_000000_003357_gtFine_labelIds.png b/datasets/cityscapes/test_label/frankfurt_000000_003357_gtFine_labelIds.png
new file mode 100755
index 0000000..1035e55
--- /dev/null
+++ b/datasets/cityscapes/test_label/frankfurt_000000_003357_gtFine_labelIds.png
diff --git a/datasets/cityscapes/test_label/frankfurt_000000_011810_gtFine_labelIds.png b/datasets/cityscapes/test_label/frankfurt_000000_011810_gtFine_labelIds.png
new file mode 100755
index 0000000..a86913b
--- /dev/null
+++ b/datasets/cityscapes/test_label/frankfurt_000000_011810_gtFine_labelIds.png
diff --git a/datasets/cityscapes/test_label/frankfurt_000000_012868_gtFine_labelIds.png b/datasets/cityscapes/test_label/frankfurt_000000_012868_gtFine_labelIds.png
new file mode 100755
index 0000000..fe81c83
--- /dev/null
+++ b/datasets/cityscapes/test_label/frankfurt_000000_012868_gtFine_labelIds.png
diff --git a/datasets/cityscapes/test_label/frankfurt_000001_013710_gtFine_labelIds.png b/datasets/cityscapes/test_label/frankfurt_000001_013710_gtFine_labelIds.png
new file mode 100755
index 0000000..72b4be4
--- /dev/null
+++ b/datasets/cityscapes/test_label/frankfurt_000001_013710_gtFine_labelIds.png
diff --git a/datasets/cityscapes/test_label/frankfurt_000001_015328_gtFine_labelIds.png b/datasets/cityscapes/test_label/frankfurt_000001_015328_gtFine_labelIds.png
new file mode 100755
index 0000000..afefb6b
--- /dev/null
+++ b/datasets/cityscapes/test_label/frankfurt_000001_015328_gtFine_labelIds.png
diff --git a/datasets/cityscapes/test_label/frankfurt_000001_023769_gtFine_labelIds.png b/datasets/cityscapes/test_label/frankfurt_000001_023769_gtFine_labelIds.png
new file mode 100755
index 0000000..f3af9df
--- /dev/null
+++ b/datasets/cityscapes/test_label/frankfurt_000001_023769_gtFine_labelIds.png
diff --git a/datasets/cityscapes/test_label/frankfurt_000001_028335_gtFine_labelIds.png b/datasets/cityscapes/test_label/frankfurt_000001_028335_gtFine_labelIds.png
new file mode 100755
index 0000000..5e65e3e
--- /dev/null
+++ b/datasets/cityscapes/test_label/frankfurt_000001_028335_gtFine_labelIds.png
diff --git a/datasets/cityscapes/test_label/frankfurt_000001_032711_gtFine_labelIds.png b/datasets/cityscapes/test_label/frankfurt_000001_032711_gtFine_labelIds.png
new file mode 100755
index 0000000..ba07b73
--- /dev/null
+++ b/datasets/cityscapes/test_label/frankfurt_000001_032711_gtFine_labelIds.png
diff --git a/datasets/cityscapes/test_label/frankfurt_000001_033655_gtFine_labelIds.png b/datasets/cityscapes/test_label/frankfurt_000001_033655_gtFine_labelIds.png
new file mode 100755
index 0000000..77f519c
--- /dev/null
+++ b/datasets/cityscapes/test_label/frankfurt_000001_033655_gtFine_labelIds.png
diff --git a/datasets/cityscapes/test_label/frankfurt_000001_042733_gtFine_labelIds.png b/datasets/cityscapes/test_label/frankfurt_000001_042733_gtFine_labelIds.png
new file mode 100755
index 0000000..ba08f1d
--- /dev/null
+++ b/datasets/cityscapes/test_label/frankfurt_000001_042733_gtFine_labelIds.png
diff --git a/datasets/cityscapes/test_label/frankfurt_000001_047552_gtFine_labelIds.png b/datasets/cityscapes/test_label/frankfurt_000001_047552_gtFine_labelIds.png
new file mode 100755
index 0000000..5dff09a
--- /dev/null
+++ b/datasets/cityscapes/test_label/frankfurt_000001_047552_gtFine_labelIds.png
diff --git a/datasets/cityscapes/test_label/frankfurt_000001_054640_gtFine_labelIds.png b/datasets/cityscapes/test_label/frankfurt_000001_054640_gtFine_labelIds.png
new file mode 100755
index 0000000..cb2ab2b
--- /dev/null
+++ b/datasets/cityscapes/test_label/frankfurt_000001_054640_gtFine_labelIds.png
diff --git a/datasets/cityscapes/test_label/frankfurt_000001_055387_gtFine_labelIds.png b/datasets/cityscapes/test_label/frankfurt_000001_055387_gtFine_labelIds.png
new file mode 100755
index 0000000..b00ef7e
--- /dev/null
+++ b/datasets/cityscapes/test_label/frankfurt_000001_055387_gtFine_labelIds.png
diff --git a/datasets/cityscapes/train_img/aachen_000000_000019_leftImg8bit.png b/datasets/cityscapes/train_img/aachen_000000_000019_leftImg8bit.png
new file mode 100755
index 0000000..0e6867e
--- /dev/null
+++ b/datasets/cityscapes/train_img/aachen_000000_000019_leftImg8bit.png
diff --git a/datasets/cityscapes/train_img/aachen_000001_000019_leftImg8bit.png b/datasets/cityscapes/train_img/aachen_000001_000019_leftImg8bit.png
new file mode 100755
index 0000000..d5a96ce
--- /dev/null
+++ b/datasets/cityscapes/train_img/aachen_000001_000019_leftImg8bit.png
diff --git a/datasets/cityscapes/train_img/aachen_000002_000019_leftImg8bit.png b/datasets/cityscapes/train_img/aachen_000002_000019_leftImg8bit.png
new file mode 100755
index 0000000..10ce563
--- /dev/null
+++ b/datasets/cityscapes/train_img/aachen_000002_000019_leftImg8bit.png
diff --git a/datasets/cityscapes/train_img/aachen_000003_000019_leftImg8bit.png b/datasets/cityscapes/train_img/aachen_000003_000019_leftImg8bit.png
new file mode 100755
index 0000000..3027fe1
--- /dev/null
+++ b/datasets/cityscapes/train_img/aachen_000003_000019_leftImg8bit.png
diff --git a/datasets/cityscapes/train_img/aachen_000004_000019_leftImg8bit.png b/datasets/cityscapes/train_img/aachen_000004_000019_leftImg8bit.png
new file mode 100755
index 0000000..26945fc
--- /dev/null
+++ b/datasets/cityscapes/train_img/aachen_000004_000019_leftImg8bit.png
diff --git a/datasets/cityscapes/train_inst/aachen_000000_000019_gtFine_instanceIds.png b/datasets/cityscapes/train_inst/aachen_000000_000019_gtFine_instanceIds.png
new file mode 100755
index 0000000..f4ee222
--- /dev/null
+++ b/datasets/cityscapes/train_inst/aachen_000000_000019_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/train_inst/aachen_000001_000019_gtFine_instanceIds.png b/datasets/cityscapes/train_inst/aachen_000001_000019_gtFine_instanceIds.png
new file mode 100755
index 0000000..dd69137
--- /dev/null
+++ b/datasets/cityscapes/train_inst/aachen_000001_000019_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/train_inst/aachen_000002_000019_gtFine_instanceIds.png b/datasets/cityscapes/train_inst/aachen_000002_000019_gtFine_instanceIds.png
new file mode 100755
index 0000000..bdad5e3
--- /dev/null
+++ b/datasets/cityscapes/train_inst/aachen_000002_000019_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/train_inst/aachen_000003_000019_gtFine_instanceIds.png b/datasets/cityscapes/train_inst/aachen_000003_000019_gtFine_instanceIds.png
new file mode 100755
index 0000000..91a035b
--- /dev/null
+++ b/datasets/cityscapes/train_inst/aachen_000003_000019_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/train_inst/aachen_000004_000019_gtFine_instanceIds.png b/datasets/cityscapes/train_inst/aachen_000004_000019_gtFine_instanceIds.png
new file mode 100755
index 0000000..0f5fc70
--- /dev/null
+++ b/datasets/cityscapes/train_inst/aachen_000004_000019_gtFine_instanceIds.png
diff --git a/datasets/cityscapes/train_label/aachen_000000_000019_gtFine_labelIds.png b/datasets/cityscapes/train_label/aachen_000000_000019_gtFine_labelIds.png
new file mode 100755
index 0000000..eed7ee6
--- /dev/null
+++ b/datasets/cityscapes/train_label/aachen_000000_000019_gtFine_labelIds.png
diff --git a/datasets/cityscapes/train_label/aachen_000001_000019_gtFine_labelIds.png b/datasets/cityscapes/train_label/aachen_000001_000019_gtFine_labelIds.png
new file mode 100755
index 0000000..e9c25ee
--- /dev/null
+++ b/datasets/cityscapes/train_label/aachen_000001_000019_gtFine_labelIds.png
diff --git a/datasets/cityscapes/train_label/aachen_000002_000019_gtFine_labelIds.png b/datasets/cityscapes/train_label/aachen_000002_000019_gtFine_labelIds.png
new file mode 100755
index 0000000..c96ab17
--- /dev/null
+++ b/datasets/cityscapes/train_label/aachen_000002_000019_gtFine_labelIds.png
diff --git a/datasets/cityscapes/train_label/aachen_000003_000019_gtFine_labelIds.png b/datasets/cityscapes/train_label/aachen_000003_000019_gtFine_labelIds.png
new file mode 100755
index 0000000..da05594
--- /dev/null
+++ b/datasets/cityscapes/train_label/aachen_000003_000019_gtFine_labelIds.png
diff --git a/datasets/cityscapes/train_label/aachen_000004_000019_gtFine_labelIds.png b/datasets/cityscapes/train_label/aachen_000004_000019_gtFine_labelIds.png
new file mode 100755
index 0000000..bb30bd9
--- /dev/null
+++ b/datasets/cityscapes/train_label/aachen_000004_000019_gtFine_labelIds.png
diff --git a/encode_features.py b/encode_features.py
new file mode 100755
index 0000000..0e97da8
--- /dev/null
+++ b/encode_features.py
@@ -0,0 +1,57 @@
+### Copyright (C) 2017 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 options.train_options import TrainOptions

+from data.data_loader import CreateDataLoader

+from models.models import create_model

+import numpy as np

+import os, time

+import util.util as util

+from torch.autograd import Variable

+

+opt = TrainOptions().parse()

+opt.nThreads = 1

+opt.batchSize = 1 

+opt.serial_batches = True 

+opt.no_flip = True

+opt.instance_feat = True

+

+name = 'features'

+save_path = os.path.join(opt.checkpoints_dir, opt.name)

+

+############ Initialize #########

+data_loader = CreateDataLoader(opt)

+dataset = data_loader.load_data()

+dataset_size = len(data_loader)

+model = create_model(opt)

+

+########### Encode features ###########

+reencode = True

+if reencode:

+	features = {}

+	for label in range(opt.label_nc):

+		features[label] = np.zeros((0, opt.feat_num+1))

+	for i, data in enumerate(dataset):    

+	    feat = model.module.encode_features(data['image'], data['inst'])

+	    for label in range(opt.label_nc):

+	    	features[label] = np.append(features[label], feat[label], axis=0) 

+	        

+	    print('%d / %d images' % (i+1, dataset_size))    

+	save_name = os.path.join(save_path, name + '.npy')

+	np.save(save_name, features)

+

+############## Clustering ###########

+n_clusters = opt.n_clusters

+load_name = os.path.join(save_path, name + '.npy')

+features = np.load(load_name).item()

+from sklearn.cluster import KMeans

+centers = {}

+for label in range(opt.label_nc):

+	feat = features[label]

+	feat = feat[feat[:,-1] > 0.5, :-1]		

+	if feat.shape[0]:

+		n_clusters = min(feat.shape[0], opt.n_clusters)

+		kmeans = KMeans(n_clusters=n_clusters, random_state=0).fit(feat)

+		centers[label] = kmeans.cluster_centers_

+save_name = os.path.join(save_path, name + '_clustered_%03d.npy' % opt.n_clusters)

+np.save(save_name, centers)

+print('saving to %s' % save_name)
\ No newline at end of file
diff --git a/models/__init__.py b/models/__init__.py
new file mode 100755
index 0000000..e69de29
--- /dev/null
+++ b/models/__init__.py
diff --git a/models/base_model.py b/models/base_model.py
new file mode 100755
index 0000000..d3879d0
--- /dev/null
+++ b/models/base_model.py
@@ -0,0 +1,86 @@
+### Copyright (C) 2017 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
+
+class BaseModel(torch.nn.Module):
+    def name(self):
+        return 'BaseModel'
+
+    def initialize(self, opt):
+        self.opt = opt
+        self.gpu_ids = opt.gpu_ids
+        self.isTrain = opt.isTrain
+        self.Tensor = torch.cuda.FloatTensor if self.gpu_ids else torch.Tensor
+        self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
+
+    def set_input(self, input):
+        self.input = input
+
+    def forward(self):
+        pass
+
+    # used in test time, no backprop
+    def test(self):
+        pass
+
+    def get_image_paths(self):
+        pass
+
+    def optimize_parameters(self):
+        pass
+
+    def get_current_visuals(self):
+        return self.input
+
+    def get_current_errors(self):
+        return {}
+
+    def save(self, label):
+        pass
+
+    # helper saving function that can be used by subclasses
+    def save_network(self, network, network_label, epoch_label, gpu_ids):
+        save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
+        save_path = os.path.join(self.save_dir, save_filename)
+        torch.save(network.cpu().state_dict(), save_path)
+        if len(gpu_ids) and torch.cuda.is_available():
+            network.cuda()
+
+    # helper loading function that can be used by subclasses
+    def load_network(self, network, network_label, epoch_label, save_dir=''):        
+        save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
+        if not save_dir:
+            save_dir = self.save_dir
+        save_path = os.path.join(save_dir, save_filename)        
+        if not os.path.isfile(save_path):
+            print('%s not exists yet!' % save_path)
+            if network_label == 'G':
+                raise('Generator must exist!')
+        else:
+            #network.load_state_dict(torch.load(save_path))
+            try:
+                network.load_state_dict(torch.load(save_path))
+            except:   
+                pretrained_dict = torch.load(save_path)                
+                model_dict = network.state_dict()
+                try:
+                    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}                    
+                    network.load_state_dict(pretrained_dict)
+                    print('Pretrained network %s has excessive layers; Only loading layers that are used' % network_label)
+                except:
+                    print('Pretrained network %s has fewer layers; The following are not initialized:' % network_label)
+                    from sets import Set
+                    not_initialized = Set()
+                    for k, v in pretrained_dict.items():                      
+                        if v.size() == model_dict[k].size():
+                            model_dict[k] = v
+
+                    for k, v in model_dict.items():
+                        if k not in pretrained_dict or v.size() != pretrained_dict[k].size():
+                            not_initialized.add(k.split('.')[0])                            
+                    print(sorted(not_initialized))
+                    network.load_state_dict(model_dict)                  
+
+    def update_learning_rate():
+        pass
diff --git a/models/models.py b/models/models.py
new file mode 100755
index 0000000..351483c
--- /dev/null
+++ b/models/models.py
@@ -0,0 +1,14 @@
+### Copyright (C) 2017 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
+
+def create_model(opt):
+    from .pix2pixHD_model import Pix2PixHDModel
+    model = Pix2PixHDModel()    
+    model.initialize(opt)
+    print("model [%s] was created" % (model.name()))
+
+    if opt.isTrain and len(opt.gpu_ids):
+        model = torch.nn.DataParallel(model, device_ids=opt.gpu_ids)
+
+    return model
diff --git a/models/networks.py b/models/networks.py
new file mode 100755
index 0000000..a673a56
--- /dev/null
+++ b/models/networks.py
@@ -0,0 +1,421 @@
+### Copyright (C) 2017 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
+from torch.nn import init
+import functools
+from torch.autograd import Variable
+import numpy as np
+import math
+import torch.nn.functional as F
+import copy
+
+###############################################################################
+# Functions
+###############################################################################
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        m.weight.data.normal_(0.0, 0.02)
+    elif classname.find('BatchNorm2d') != -1:
+        m.weight.data.normal_(1.0, 0.02)
+        m.bias.data.fill_(0)
+
+def get_norm_layer(norm_type='instance'):
+    if norm_type == 'batch':
+        norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
+    elif norm_type == 'instance':
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False)
+    else:
+        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
+    return norm_layer
+
+def define_G(input_nc, output_nc, ngf, netG, n_downsample_global=3, n_blocks_global=9, n_local_enhancers=1, 
+             n_blocks_local=3, norm='instance', gpu_ids=[]):    
+    norm_layer = get_norm_layer(norm_type=norm)     
+    if netG == 'global':    
+        netG = GlobalGenerator(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, norm_layer)       
+    elif netG == 'local':        
+        netG = LocalEnhancer(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, 
+                                  n_local_enhancers, n_blocks_local, norm_layer)
+    elif netG == 'encoder':
+        netG = Encoder(input_nc, output_nc, ngf, n_downsample_global, norm_layer)
+    else:
+        raise('generator not implemented!')
+    print(netG)
+    if len(gpu_ids) > 0:
+        assert(torch.cuda.is_available())   
+        netG.cuda(device_id=gpu_ids[0])
+    netG.apply(weights_init)
+    return netG
+
+def define_D(input_nc, ndf, n_layers_D, norm='instance', use_sigmoid=False, num_D=1, getIntermFeat=False, gpu_ids=[]):        
+    norm_layer = get_norm_layer(norm_type=norm)   
+    netD = MultiscaleDiscriminator(input_nc, ndf, n_layers_D, norm_layer, use_sigmoid, num_D, getIntermFeat)   
+    print(netD)
+    if len(gpu_ids) > 0:
+        assert(torch.cuda.is_available())
+        netD.cuda(device_id=gpu_ids[0])
+    netD.apply(weights_init)
+    return netD
+
+def print_network(net):
+    if isinstance(net, list):
+        net = net[0]
+    num_params = 0
+    for param in net.parameters():
+        num_params += param.numel()
+    print(net)
+    print('Total number of parameters: %d' % num_params)
+
+##############################################################################
+# Losses
+##############################################################################
+class GANLoss(nn.Module):
+    def __init__(self, use_lsgan=True, target_real_label=1.0, target_fake_label=0.0,
+                 tensor=torch.FloatTensor):
+        super(GANLoss, self).__init__()
+        self.real_label = target_real_label
+        self.fake_label = target_fake_label
+        self.real_label_var = None
+        self.fake_label_var = None
+        self.Tensor = tensor
+        if use_lsgan:
+            self.loss = nn.MSELoss()
+        else:
+            self.loss = nn.BCELoss()
+
+    def get_target_tensor(self, input, target_is_real):
+        target_tensor = None
+        if target_is_real:
+            create_label = ((self.real_label_var is None) or
+                            (self.real_label_var.numel() != input.numel()))
+            if create_label:
+                real_tensor = self.Tensor(input.size()).fill_(self.real_label)
+                self.real_label_var = Variable(real_tensor, requires_grad=False)
+            target_tensor = self.real_label_var
+        else:
+            create_label = ((self.fake_label_var is None) or
+                            (self.fake_label_var.numel() != input.numel()))
+            if create_label:
+                fake_tensor = self.Tensor(input.size()).fill_(self.fake_label)
+                self.fake_label_var = Variable(fake_tensor, requires_grad=False)
+            target_tensor = self.fake_label_var
+        return target_tensor
+
+    def __call__(self, input, target_is_real):
+        if isinstance(input[0], list):
+            loss = 0
+            for input_i in input:
+                pred = input_i[-1]
+                target_tensor = self.get_target_tensor(pred, target_is_real)
+                loss += self.loss(pred, target_tensor)
+            return loss
+        else:            
+            target_tensor = self.get_target_tensor(input[-1], target_is_real)
+            return self.loss(input[-1], target_tensor)
+
+class VGGLoss(nn.Module):
+    def __init__(self, gpu_ids):
+        super(VGGLoss, self).__init__()        
+        self.vgg = Vgg19().cuda()
+        self.criterion = nn.L1Loss()
+        self.weights = [1.0/32, 1.0/16, 1.0/8, 1.0/4, 1.0]        
+
+    def forward(self, x, y):              
+        x_vgg, y_vgg = self.vgg(x), self.vgg(y)
+        loss = 0
+        for i in range(len(x_vgg)):
+            loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())        
+        return loss
+
+##############################################################################
+# Generator
+##############################################################################
+class LocalEnhancer(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=32, n_downsample_global=3, n_blocks_global=9, 
+                 n_local_enhancers=1, n_blocks_local=3, norm_layer=nn.BatchNorm2d, padding_type='reflect'):        
+        super(LocalEnhancer, self).__init__()
+        self.n_local_enhancers = n_local_enhancers
+        
+        ###### global generator model #####           
+        ngf_global = ngf * (2**n_local_enhancers)
+        model_global = GlobalGenerator(input_nc, output_nc, ngf_global, n_downsample_global, n_blocks_global, norm_layer).model        
+        model_global = [model_global[i] for i in range(len(model_global)-3)] # get rid of final convolution layers        
+        self.model = nn.Sequential(*model_global)                
+
+        ###### local enhancer layers #####
+        for n in range(1, n_local_enhancers+1):
+            ### downsample            
+            ngf_global = ngf * (2**(n_local_enhancers-n))
+            model_downsample = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf_global, kernel_size=7, padding=0), 
+                                norm_layer(ngf_global), nn.ReLU(True),
+                                nn.Conv2d(ngf_global, ngf_global * 2, kernel_size=3, stride=2, padding=1), 
+                                norm_layer(ngf_global * 2), nn.ReLU(True)]
+            ### residual blocks
+            model_upsample = []
+            for i in range(n_blocks_local):
+                model_upsample += [ResnetBlock(ngf_global * 2, padding_type=padding_type, norm_layer=norm_layer)]
+
+            ### upsample
+            model_upsample += [nn.ConvTranspose2d(ngf_global * 2, ngf_global, kernel_size=3, stride=2, padding=1, output_padding=1), 
+                               norm_layer(ngf_global), nn.ReLU(True)]      
+
+            ### final convolution
+            if n == n_local_enhancers:                
+                model_upsample += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]                       
+            
+            setattr(self, 'model'+str(n)+'_1', nn.Sequential(*model_downsample))
+            setattr(self, 'model'+str(n)+'_2', nn.Sequential(*model_upsample))                  
+        
+        self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
+
+    def forward(self, input): 
+        ### create input pyramid
+        input_downsampled = [input]
+        for i in range(self.n_local_enhancers):
+            input_downsampled.append(self.downsample(input_downsampled[-1]))
+
+        ### output at coarest level
+        output_prev = self.model(input_downsampled[-1])        
+        ### build up one layer at a time
+        for n_local_enhancers in range(1, self.n_local_enhancers+1):
+            model_downsample = getattr(self, 'model'+str(n_local_enhancers)+'_1')
+            model_upsample = getattr(self, 'model'+str(n_local_enhancers)+'_2')            
+            input_i = input_downsampled[self.n_local_enhancers-n_local_enhancers]            
+            output_prev = model_upsample(model_downsample(input_i) + output_prev)
+        return output_prev
+
+class GlobalGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=64, n_downsampling=3, n_blocks=9, norm_layer=nn.BatchNorm2d, 
+                 padding_type='reflect'):
+        assert(n_blocks >= 0)
+        super(GlobalGenerator, self).__init__()        
+        activation = nn.ReLU(True)        
+
+        model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
+        ### downsample
+        for i in range(n_downsampling):
+            mult = 2**i
+            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
+                      norm_layer(ngf * mult * 2), activation]
+
+        ### resnet blocks
+        mult = 2**n_downsampling
+        for i in range(n_blocks):
+            model += [ResnetBlock(ngf * mult, padding_type=padding_type, activation=activation, norm_layer=norm_layer)]
+        
+        ### upsample         
+        for i in range(n_downsampling):
+            mult = 2**(n_downsampling - i)
+            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1),
+                       norm_layer(int(ngf * mult / 2)), activation]
+        model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]        
+        self.model = nn.Sequential(*model)
+            
+    def forward(self, input):
+        return self.model(input)             
+        
+# Define a resnet block
+class ResnetBlock(nn.Module):
+    def __init__(self, dim, padding_type, norm_layer, activation=nn.ReLU(True), use_dropout=False):
+        super(ResnetBlock, self).__init__()
+        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, activation, use_dropout)
+
+    def build_conv_block(self, dim, padding_type, norm_layer, activation, use_dropout):
+        conv_block = []
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p),
+                       norm_layer(dim),
+                       activation]
+        if use_dropout:
+            conv_block += [nn.Dropout(0.5)]
+
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p),
+                       norm_layer(dim)]
+
+        return nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        out = x + self.conv_block(x)
+        return out
+
+class Encoder(nn.Module):
+    def __init__(self, input_nc, output_nc, ngf=32, n_downsampling=4, norm_layer=nn.BatchNorm2d):
+        super(Encoder, self).__init__()        
+        self.output_nc = output_nc        
+
+        model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), 
+                 norm_layer(ngf), nn.ReLU(True)]             
+        ### downsample
+        for i in range(n_downsampling):
+            mult = 2**i
+            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
+                      norm_layer(ngf * mult * 2), nn.ReLU(True)]
+
+        ### upsample         
+        for i in range(n_downsampling):
+            mult = 2**(n_downsampling - i)
+            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1),
+                       norm_layer(int(ngf * mult / 2)), nn.ReLU(True)]        
+
+        model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
+        self.model = nn.Sequential(*model) 
+
+    def forward(self, input, inst):
+        outputs = self.model(input)
+
+        # instance-wise average pooling
+        outputs_mean = outputs.clone()        
+        inst_list = np.unique(inst.cpu().numpy().astype(int))        
+        for i in inst_list:
+            indices = (inst == i).nonzero() # n x 4            
+            for j in range(self.output_nc):
+                output_ins = outputs[indices[:,0], indices[:,1] + j, indices[:,2], indices[:,3]]                    
+                mean_feat = torch.mean(output_ins).expand_as(output_ins)                                        
+                outputs_mean[indices[:,0], indices[:,1] + j, indices[:,2], indices[:,3]] = mean_feat                        
+        return outputs_mean
+
+class MultiscaleDiscriminator(nn.Module):
+    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, 
+                 use_sigmoid=False, num_D=3, getIntermFeat=False):
+        super(MultiscaleDiscriminator, self).__init__()
+        self.num_D = num_D
+        self.n_layers = n_layers
+        self.getIntermFeat = getIntermFeat
+     
+        for i in range(num_D):
+            netD = NLayerDiscriminator(input_nc, ndf, n_layers, norm_layer, use_sigmoid, getIntermFeat)
+            if getIntermFeat:                                
+                for j in range(n_layers+2):
+                    setattr(self, 'scale'+str(i)+'_layer'+str(j), getattr(netD, 'model'+str(j)))                                   
+            else:
+                setattr(self, 'layer'+str(i), netD.model)
+
+        self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
+
+    def singleD_forward(self, model, input):
+        if self.getIntermFeat:
+            result = [input]
+            for i in range(len(model)):
+                result.append(model[i](result[-1]))
+            return result[1:]
+        else:
+            return [model(input)]
+
+    def forward(self, input):        
+        num_D = self.num_D
+        result = []
+        input_downsampled = input
+        for i in range(num_D):
+            if self.getIntermFeat:
+                model = [getattr(self, 'scale'+str(num_D-1-i)+'_layer'+str(j)) for j in range(self.n_layers+2)]
+            else:
+                model = getattr(self, 'layer'+str(num_D-1-i))
+            result.append(self.singleD_forward(model, input_downsampled))
+            if i != (num_D-1):
+                input_downsampled = self.downsample(input_downsampled)
+        return result
+        
+# Defines the PatchGAN discriminator with the specified arguments.
+class NLayerDiscriminator(nn.Module):
+    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_sigmoid=False, getIntermFeat=False):
+        super(NLayerDiscriminator, self).__init__()
+        self.getIntermFeat = getIntermFeat
+        self.n_layers = n_layers
+
+        kw = 4
+        padw = int(np.ceil((kw-1.0)/2))
+        sequence = [[nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]]
+
+        nf = ndf
+        for n in range(1, n_layers):
+            nf_prev = nf
+            nf = min(nf * 2, 512)
+            sequence += [[
+                nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=2, padding=padw),
+                norm_layer(nf), nn.LeakyReLU(0.2, True)
+            ]]
+
+        nf_prev = nf
+        nf = min(nf * 2, 512)
+        sequence += [[
+            nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=1, padding=padw),
+            norm_layer(nf),
+            nn.LeakyReLU(0.2, True)
+        ]]
+
+        sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)]]
+
+        if use_sigmoid:
+            sequence += [nn.Sigmoid()]
+
+        if getIntermFeat:
+            for n in range(len(sequence)):
+                setattr(self, 'model'+str(n), nn.Sequential(*sequence[n]))
+        else:
+            sequence_stream = []
+            for n in range(len(sequence)):
+                sequence_stream += sequence[n]
+            self.model = nn.Sequential(*sequence_stream)
+
+    def forward(self, input):
+        if self.getIntermFeat:
+            res = [input]
+            for n in range(self.n_layers+2):
+                model = getattr(self, 'model'+str(n))
+                res.append(model(res[-1]))
+            return res[1:]
+        else:
+            return self.model(input)        
+
+from torchvision import models
+class Vgg19(torch.nn.Module):
+    def __init__(self, requires_grad=False):
+        super(Vgg19, self).__init__()
+        vgg_pretrained_features = 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
+
+    def forward(self, X):
+        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
diff --git a/models/pix2pixHD_model.py b/models/pix2pixHD_model.py
new file mode 100755
index 0000000..ba44e53
--- /dev/null
+++ b/models/pix2pixHD_model.py
@@ -0,0 +1,260 @@
+### Copyright (C) 2017 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
+import torch
+import os
+from collections import OrderedDict
+from torch.autograd import Variable
+import util.util as util
+from util.image_pool import ImagePool
+from .base_model import BaseModel
+from . import networks
+
+class Pix2PixHDModel(BaseModel):
+    def name(self):
+        return 'Pix2PixHDModel'
+
+    def initialize(self, opt):
+        BaseModel.initialize(self, opt)
+        if opt.resize_or_crop != 'none': # when training at full res this causes OOM
+            torch.backends.cudnn.benchmark = True
+        self.isTrain = opt.isTrain
+        self.use_features = opt.instance_feat or opt.label_feat
+        self.gen_features = self.use_features and not self.opt.load_features
+
+        ##### define networks        
+        # Generator network
+        netG_input_nc = opt.label_nc
+        if not opt.no_instance:
+            netG_input_nc += 1
+        if self.use_features:
+            netG_input_nc += opt.feat_num                  
+        self.netG = networks.define_G(netG_input_nc, opt.output_nc, opt.ngf, opt.netG, 
+                                      opt.n_downsample_global, opt.n_blocks_global, opt.n_local_enhancers, 
+                                      opt.n_blocks_local, opt.norm, gpu_ids=self.gpu_ids)        
+
+        # Discriminator network
+        if self.isTrain:
+            use_sigmoid = opt.no_lsgan
+            netD_input_nc = opt.label_nc + opt.output_nc
+            if not opt.no_instance:
+                netD_input_nc += 1
+            self.netD = networks.define_D(netD_input_nc, opt.ndf, opt.n_layers_D, opt.norm, use_sigmoid, 
+                                          opt.num_D, not opt.no_ganFeat_loss, gpu_ids=self.gpu_ids)
+
+        ### Encoder network
+        if self.gen_features:          
+            self.netE = networks.define_G(opt.output_nc, opt.feat_num, opt.nef, 'encoder', 
+                                          opt.n_downsample_E, norm=opt.norm, gpu_ids=self.gpu_ids)  
+            
+        print('---------- Networks initialized -------------')
+
+        # load networks
+        if not self.isTrain or opt.continue_train or opt.load_pretrain:
+            pretrained_path = '' if not self.isTrain else opt.load_pretrain
+            self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path)            
+            if self.isTrain:
+                self.load_network(self.netD, 'D', opt.which_epoch, pretrained_path)  
+            if self.gen_features:
+                self.load_network(self.netE, 'E', opt.which_epoch, pretrained_path)              
+
+        # set loss functions and optimizers
+        if self.isTrain:
+            if opt.pool_size > 0 and (len(self.gpu_ids)) > 1:
+                raise NotImplementedError("Fake Pool Not Implemented for MultiGPU")
+            self.fake_pool = ImagePool(opt.pool_size)
+            self.old_lr = opt.lr
+
+            # define loss functions
+            self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan, tensor=self.Tensor)   
+            self.criterionFeat = torch.nn.L1Loss()
+            if not opt.no_vgg_loss:             
+                self.criterionVGG = networks.VGGLoss(self.gpu_ids)
+        
+            # Names so we can breakout loss
+            self.loss_names = ['G_GAN', 'G_GAN_Feat', 'G_VGG', 'D_real', 'D_fake']
+
+            # initialize optimizers
+            # optimizer G
+            if opt.niter_fix_global > 0:
+                print('------------- Only training the local enhancer network (for %d epochs) ------------' % opt.niter_fix_global)
+                params_dict = dict(self.netG.named_parameters())
+                params = []
+                for key, value in params_dict.items():       
+                    if key.startswith('model' + str(opt.n_local_enhancers)):
+                        params += [{'params':[value],'lr':opt.lr}]
+                    else:
+                        params += [{'params':[value],'lr':0.0}]                            
+            else:
+                params = list(self.netG.parameters())
+            if self.gen_features:              
+                params += list(self.netE.parameters())         
+            self.optimizer_G = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))                            
+
+            # optimizer D                        
+            params = list(self.netD.parameters())    
+            self.optimizer_D = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
+
+    def encode_input(self, label_map, inst_map=None, real_image=None, feat_map=None, infer=False):     
+        # create one-hot vector for label map 
+        size = label_map.size()
+        oneHot_size = (size[0], self.opt.label_nc, size[2], size[3])
+        input_label = torch.cuda.FloatTensor(torch.Size(oneHot_size)).zero_()
+        input_label = input_label.scatter_(1, label_map.data.long().cuda(), 1.0)
+
+        # get edges from instance map
+        if not self.opt.no_instance:
+            inst_map = inst_map.data.cuda()
+            edge_map = self.get_edges(inst_map)
+            input_label = torch.cat((input_label, edge_map), dim=1) 
+        input_label = Variable(input_label, volatile=infer)
+
+        # real images for training
+        if real_image is not None:
+            real_image = Variable(real_image.data.cuda())
+
+        # instance map for feature encoding
+        if self.use_features:
+            # get precomputed feature maps
+            if self.opt.load_features:
+                feat_map = Variable(feat_map.data.cuda())
+
+        return input_label, inst_map, real_image, feat_map
+
+    def discriminate(self, input_label, test_image, use_pool=False):
+        input_concat = torch.cat((input_label, test_image.detach()), dim=1)
+        if use_pool:            
+            fake_query = self.fake_pool.query(input_concat)
+            return self.netD.forward(fake_query)
+        else:
+            return self.netD.forward(input_concat)
+
+    def forward(self, label, inst, image, feat, infer=False):
+        # Encode Inputs
+        input_label, inst_map, real_image, feat_map = self.encode_input(label, inst, image, feat)  
+
+        # Fake Generation
+        if self.use_features:
+            if not self.opt.load_features:
+                feat_map = self.netE.forward(real_image, inst_map)                     
+            input_concat = torch.cat((input_label, feat_map), dim=1)                        
+        else:
+            input_concat = input_label
+        fake_image = self.netG.forward(input_concat)
+
+        # Fake Detection and Loss
+        pred_fake_pool = self.discriminate(input_label, fake_image, use_pool=True)
+        loss_D_fake = self.criterionGAN(pred_fake_pool, False)        
+
+        # Real Detection and Loss        
+        pred_real = self.discriminate(input_label, real_image)
+        loss_D_real = self.criterionGAN(pred_real, True)
+
+        # GAN loss (Fake Passability Loss)        
+        pred_fake = self.netD.forward(torch.cat((input_label, fake_image), dim=1))        
+        loss_G_GAN = self.criterionGAN(pred_fake, True)               
+        
+        # GAN feature matching loss
+        loss_G_GAN_Feat = 0
+        if not self.opt.no_ganFeat_loss:
+            feat_weights = 4.0 / (self.opt.n_layers_D + 1)
+            D_weights = 1.0 / self.opt.num_D
+            for i in range(self.opt.num_D):
+                for j in range(len(pred_fake[i])-1):
+                    loss_G_GAN_Feat += D_weights * feat_weights * \
+                        self.criterionFeat(pred_fake[i][j], pred_real[i][j].detach()) * self.opt.lambda_feat
+                   
+        # VGG feature matching loss
+        loss_G_VGG = 0
+        if not self.opt.no_vgg_loss:
+            loss_G_VGG = self.criterionVGG(fake_image, real_image) * self.opt.lambda_feat
+        
+        # Only return the fake_B image if necessary to save BW
+        return [ [ loss_G_GAN, loss_G_GAN_Feat, loss_G_VGG, loss_D_real, loss_D_fake ], None if not infer else fake_image ]
+
+    def inference(self, label, inst):
+        # Encode Inputs        
+        input_label, inst_map, _, _ = self.encode_input(Variable(label), Variable(inst), infer=True)
+
+        # Fake Generation
+        if self.use_features:       
+            # sample clusters from precomputed features             
+            feat_map = self.sample_features(inst_map)
+            input_concat = torch.cat((input_label, feat_map), dim=1)                        
+        else:
+            input_concat = input_label                
+        fake_image = self.netG.forward(input_concat)
+        return fake_image
+
+    def sample_features(self, inst): 
+        # read precomputed feature clusters 
+        cluster_path = os.path.join(self.opt.checkpoints_dir, self.opt.name, self.opt.cluster_path)        
+        features_clustered = np.load(cluster_path).item()
+
+        # randomly sample from the feature clusters
+        inst_np = inst.cpu().numpy().astype(int)                                      
+        feat_map = torch.cuda.FloatTensor(1, self.opt.feat_num, inst.size()[2], inst.size()[3])                   
+        for i in np.unique(inst_np):    
+            label = i if i < 1000 else i//1000
+            if label in features_clustered:
+                feat = features_clustered[label]
+                cluster_idx = np.random.randint(0, feat.shape[0]) 
+                                            
+                idx = (inst == i).nonzero()
+                for k in range(self.opt.feat_num):                                    
+                    feat_map[idx[:,0], idx[:,1] + k, idx[:,2], idx[:,3]] = feat[cluster_idx, k] 
+        return feat_map
+
+    def encode_features(self, image, inst):
+        image = Variable(image.cuda(), volatile=True)
+        feat_num = self.opt.feat_num
+        h, w = inst.size()[2], inst.size()[3]
+        block_num = 32
+        feat_map = self.netE.forward(image, inst.cuda())
+        inst_np = inst.cpu().numpy().astype(int)
+        feature = {}
+        for i in range(self.opt.label_nc):
+            feature[i] = np.zeros((0, feat_num+1))
+        for i in np.unique(inst_np):
+            label = i if i < 1000 else i//1000
+            idx = (inst == i).nonzero()
+            num = idx.size()[0]
+            idx = idx[num//2,:]
+            val = np.zeros((1, feat_num+1))                        
+            for k in range(feat_num):
+                val[0, k] = feat_map[idx[0], idx[1] + k, idx[2], idx[3]].data[0]            
+            val[0, feat_num] = float(num) / (h * w // block_num)
+            feature[label] = np.append(feature[label], val, axis=0)
+        return feature
+
+    def get_edges(self, t):
+        edge = torch.cuda.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 save(self, which_epoch):
+        self.save_network(self.netG, 'G', which_epoch, self.gpu_ids)
+        self.save_network(self.netD, 'D', which_epoch, self.gpu_ids)
+        if self.gen_features:
+            self.save_network(self.netE, 'E', which_epoch, self.gpu_ids)
+
+    def update_fixed_params(self):
+        # after fixing the global generator for a number of iterations, also start finetuning it
+        params = list(self.netG.parameters())
+        if self.gen_features:
+            params += list(self.netE.parameters())           
+        self.optimizer_G = torch.optim.Adam(params, lr=self.opt.lr, betas=(self.opt.beta1, 0.999)) 
+        print('------------ Now also finetuning global generator -----------')
+
+    def update_learning_rate(self):
+        lrd = self.opt.lr / self.opt.niter_decay
+        lr = self.old_lr - lrd        
+        for param_group in self.optimizer_D.param_groups:
+            param_group['lr'] = lr
+        for param_group in self.optimizer_G.param_groups:
+            param_group['lr'] = lr
+        print('update learning rate: %f -> %f' % (self.old_lr, lr))
+        self.old_lr = lr
diff --git a/options/__init__.py b/options/__init__.py
new file mode 100755
index 0000000..e69de29
--- /dev/null
+++ b/options/__init__.py
diff --git a/options/base_options.py b/options/base_options.py
new file mode 100755
index 0000000..863c061
--- /dev/null
+++ b/options/base_options.py
@@ -0,0 +1,95 @@
+### Copyright (C) 2017 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 argparse
+import os
+from util import util
+import torch
+
+class BaseOptions():
+    def __init__(self):
+        self.parser = argparse.ArgumentParser()
+        self.initialized = False
+
+    def initialize(self):    
+        # experiment specifics
+        self.parser.add_argument('--name', type=str, default='label2city', help='name of the experiment. It decides where to store samples and models')        
+        self.parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')                       
+        self.parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
+        self.parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization')        
+        self.parser.add_argument('--use_dropout', action='store_true', help='use dropout for the generator')
+
+        # input/output sizes       
+        self.parser.add_argument('--batchSize', type=int, default=1, help='input batch size')
+        self.parser.add_argument('--loadSize', type=int, default=1024, help='scale images to this size')
+        self.parser.add_argument('--fineSize', type=int, default=512, help='then crop to this size')
+        self.parser.add_argument('--label_nc', type=int, default=35, help='# of input image channels')
+        self.parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels')
+
+        # for setting inputs
+        self.parser.add_argument('--dataroot', type=str, default='./datasets/cityscapes/') 
+        self.parser.add_argument('--resize_or_crop', type=str, default='scale_width', help='scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop]')
+        self.parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')        
+        self.parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data argumentation') 
+        self.parser.add_argument('--nThreads', default=2, type=int, help='# threads for loading data')                
+        self.parser.add_argument('--max_dataset_size', type=int, default=float("inf"), help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')
+
+        # for displays
+        self.parser.add_argument('--display_winsize', type=int, default=512,  help='display window size')
+        self.parser.add_argument('--tf_log', action='store_true', help='if specified, use tensorboard logging. Requires tensorflow installed')
+
+        # for generator
+        self.parser.add_argument('--netG', type=str, default='global', help='selects model to use for netG')
+        self.parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer')
+        self.parser.add_argument('--n_downsample_global', type=int, default=4, help='number of downsampling layers in netG') 
+        self.parser.add_argument('--n_blocks_global', type=int, default=9, help='number of residual blocks in the global generator network')
+        self.parser.add_argument('--n_blocks_local', type=int, default=3, help='number of residual blocks in the local enhancer network')
+        self.parser.add_argument('--n_local_enhancers', type=int, default=1, help='number of local enhancers to use')        
+        self.parser.add_argument('--niter_fix_global', type=int, default=0, help='number of epochs that we only train the outmost local enhancer')        
+
+        # for instance-wise features
+        self.parser.add_argument('--no_instance', action='store_true', help='if specified, do *not* add instance map as input')        
+        self.parser.add_argument('--instance_feat', action='store_true', help='if specified, add encoded instance features as input')
+        self.parser.add_argument('--label_feat', action='store_true', help='if specified, add encoded label features as input')        
+        self.parser.add_argument('--feat_num', type=int, default=3, help='vector length for encoded features')        
+        self.parser.add_argument('--load_features', action='store_true', help='if specified, load precomputed feature maps')
+        self.parser.add_argument('--n_downsample_E', type=int, default=4, help='# of downsampling layers in encoder') 
+        self.parser.add_argument('--nef', type=int, default=16, help='# of encoder filters in the first conv layer')        
+        self.parser.add_argument('--n_clusters', type=int, default=10, help='number of clusters for features')
+        
+        self.initialized = True
+
+    def parse(self, save=True):
+        if not self.initialized:
+            self.initialize()
+        self.opt = self.parser.parse_args()
+        self.opt.isTrain = self.isTrain   # train or test
+
+        str_ids = self.opt.gpu_ids.split(',')
+        self.opt.gpu_ids = []
+        for str_id in str_ids:
+            id = int(str_id)
+            if id >= 0:
+                self.opt.gpu_ids.append(id)
+        
+        # set gpu ids
+        if len(self.opt.gpu_ids) > 0:
+            torch.cuda.set_device(self.opt.gpu_ids[0])
+
+        args = vars(self.opt)
+
+        print('------------ Options -------------')
+        for k, v in sorted(args.items()):
+            print('%s: %s' % (str(k), str(v)))
+        print('-------------- End ----------------')
+
+        # save to the disk        
+        expr_dir = os.path.join(self.opt.checkpoints_dir, self.opt.name)
+        util.mkdirs(expr_dir)
+        if save and not self.opt.continue_train:
+            file_name = os.path.join(expr_dir, 'opt.txt')
+            with open(file_name, 'wt') as opt_file:
+                opt_file.write('------------ Options -------------\n')
+                for k, v in sorted(args.items()):
+                    opt_file.write('%s: %s\n' % (str(k), str(v)))
+                opt_file.write('-------------- End ----------------\n')
+        return self.opt
diff --git a/options/test_options.py b/options/test_options.py
new file mode 100755
index 0000000..aaeff53
--- /dev/null
+++ b/options/test_options.py
@@ -0,0 +1,15 @@
+### Copyright (C) 2017 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):
+        BaseOptions.initialize(self)
+        self.parser.add_argument('--ntest', type=int, default=float("inf"), help='# of test examples.')
+        self.parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.')
+        self.parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images')
+        self.parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc')
+        self.parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
+        self.parser.add_argument('--how_many', type=int, default=50, help='how many test images to run')       
+        self.parser.add_argument('--cluster_path', type=str, default='features_clustered_010.npy', help='the path for clustered results of encoded features')
+        self.isTrain = False
diff --git a/options/train_options.py b/options/train_options.py
new file mode 100755
index 0000000..9994a4a
--- /dev/null
+++ b/options/train_options.py
@@ -0,0 +1,36 @@
+### Copyright (C) 2017 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 TrainOptions(BaseOptions):
+    def initialize(self):
+        BaseOptions.initialize(self)
+        # for displays
+        self.parser.add_argument('--display_freq', type=int, default=100, help='frequency of showing training results on screen')
+        self.parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console')
+        self.parser.add_argument('--save_latest_freq', type=int, default=1000, help='frequency of saving the latest results')
+        self.parser.add_argument('--save_epoch_freq', type=int, default=10, help='frequency of saving checkpoints at the end of epochs')        
+        self.parser.add_argument('--no_html', action='store_true', help='do not save intermediate training results to [opt.checkpoints_dir]/[opt.name]/web/')
+        self.parser.add_argument('--debug', action='store_true', help='only do one epoch and displays at each iteration')
+
+        # for training
+        self.parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')
+        self.parser.add_argument('--load_pretrain', type=str, default='', help='load the pretrained model from the specified location')
+        self.parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
+        self.parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')
+        self.parser.add_argument('--niter', type=int, default=100, help='# of iter at starting learning rate')
+        self.parser.add_argument('--niter_decay', type=int, default=100, help='# of iter to linearly decay learning rate to zero')
+        self.parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam')
+        self.parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')
+
+        # for discriminators        
+        self.parser.add_argument('--num_D', type=int, default=2, help='number of discriminators to use')
+        self.parser.add_argument('--n_layers_D', type=int, default=3, help='only used if which_model_netD==n_layers')
+        self.parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in first conv layer')    
+        self.parser.add_argument('--lambda_feat', type=float, default=10.0, help='weight for feature matching loss')                
+        self.parser.add_argument('--no_ganFeat_loss', action='store_true', help='if specified, do *not* use discriminator feature matching loss')
+        self.parser.add_argument('--no_vgg_loss', action='store_true', help='if specified, do *not* use VGG feature matching loss')        
+        self.parser.add_argument('--no_lsgan', action='store_true', help='do *not* use least square GAN, if false, use vanilla GAN')
+        self.parser.add_argument('--pool_size', type=int, default=0, help='the size of image buffer that stores previously generated images')
+
+        self.isTrain = True
diff --git a/precompute_feature_maps.py b/precompute_feature_maps.py
new file mode 100755
index 0000000..a631b9c
--- /dev/null
+++ b/precompute_feature_maps.py
@@ -0,0 +1,36 @@
+### Copyright (C) 2017 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 options.train_options import TrainOptions

+from data.data_loader import CreateDataLoader

+from models.models import create_model

+import numpy as np

+import os, time

+import util.util as util

+from torch.autograd import Variable

+import torch.nn as nn

+

+opt = TrainOptions().parse()

+opt.nThreads = 1

+opt.batchSize = 1 

+opt.serial_batches = True 

+opt.no_flip = True

+opt.instance_feat = True

+

+name = 'features'

+save_path = os.path.join(opt.checkpoints_dir, opt.name)

+

+############ Initialize #########

+data_loader = CreateDataLoader(opt)

+dataset = data_loader.load_data()

+dataset_size = len(data_loader)

+model = create_model(opt)

+util.mkdirs(os.path.join(opt.dataroot, opt.phase + '_feat'))

+

+######## Save precomputed feature maps for 1024p training #######

+for i, data in enumerate(dataset):

+	print('%d / %d images' % (i+1, dataset_size)) 

+	feat_map = model.module.netE.forward(Variable(data['image'].cuda(), volatile=True), data['inst'].cuda())

+	feat_map = nn.Upsample(scale_factor=2, mode='nearest')(feat_map)

+	image_numpy = util.tensor2im(feat_map.data[0])

+	save_path = data['path'][0].replace('/train_label/', '/train_feat/')

+	util.save_image(image_numpy, save_path)
\ No newline at end of file
diff --git a/scripts/test_1024p.sh b/scripts/test_1024p.sh
new file mode 100755
index 0000000..99c1e24
--- /dev/null
+++ b/scripts/test_1024p.sh
@@ -0,0 +1,3 @@
+################################ Testing ################################

+# labels only

+python test.py --name label2city_1024p --netG local --ngf 32 --resize_or_crop none
\ No newline at end of file
diff --git a/scripts/test_1024p_feat.sh b/scripts/test_1024p_feat.sh
new file mode 100755
index 0000000..2f4ba17
--- /dev/null
+++ b/scripts/test_1024p_feat.sh
@@ -0,0 +1,5 @@
+################################ Testing ################################

+# first precompute and cluster all features

+python encode_features.py --name label2city_1024p_feat --netG local --ngf 32 --resize_or_crop none;

+# use instance-wise features

+python test.py --name label2city_1024p_feat ---netG local --ngf 32 --resize_or_crop none --instance_feat
\ No newline at end of file
diff --git a/scripts/test_512p.sh b/scripts/test_512p.sh
new file mode 100755
index 0000000..3131043
--- /dev/null
+++ b/scripts/test_512p.sh
@@ -0,0 +1,3 @@
+################################ Testing ################################

+# labels only

+python test.py --name label2city_512p
\ No newline at end of file
diff --git a/scripts/test_512p_feat.sh b/scripts/test_512p_feat.sh
new file mode 100755
index 0000000..8f25e9c
--- /dev/null
+++ b/scripts/test_512p_feat.sh
@@ -0,0 +1,5 @@
+################################ Testing ################################

+# first precompute and cluster all features

+python encode_features.py --name label2city_512p_feat;

+# use instance-wise features

+python test.py --name label2city_512p_feat --instance_feat
\ No newline at end of file
diff --git a/scripts/train_1024p_12G.sh b/scripts/train_1024p_12G.sh
new file mode 100755
index 0000000..d5ea7d7
--- /dev/null
+++ b/scripts/train_1024p_12G.sh
@@ -0,0 +1,4 @@
+############## To train images at 2048 x 1024 resolution after training 1024 x 512 resolution models #############

+##### Using GPUs with 12G memory (not tested)

+# Using labels only

+python train.py --name label2city_1024p --netG local --ngf 32 --num_D 3 --load_pretrain checkpoints/label2city_512p/ --niter_fix_global 20 --resize_or_crop crop --fineSize 1024
\ No newline at end of file
diff --git a/scripts/train_1024p_24G.sh b/scripts/train_1024p_24G.sh
new file mode 100755
index 0000000..88e58f7
--- /dev/null
+++ b/scripts/train_1024p_24G.sh
@@ -0,0 +1,4 @@
+############## To train images at 2048 x 1024 resolution after training 1024 x 512 resolution models #############

+######## Using GPUs with 24G memory

+# Using labels only

+python train.py --name label2city_1024p --netG local --ngf 32 --num_D 3 --load_pretrain checkpoints/label2city_512p/ --niter 50 --niter_decay 50 --niter_fix_global 10 --resize_or_crop none
\ No newline at end of file
diff --git a/scripts/train_1024p_feat_12G.sh b/scripts/train_1024p_feat_12G.sh
new file mode 100755
index 0000000..f8e3d61
--- /dev/null
+++ b/scripts/train_1024p_feat_12G.sh
@@ -0,0 +1,6 @@
+############## To train images at 2048 x 1024 resolution after training 1024 x 512 resolution models #############

+##### Using GPUs with 12G memory (not tested)

+# First precompute feature maps and save them

+python precompute_feature_maps.py --name label2city_512p_feat;

+# Adding instances and encoded features

+python train.py --name label2city_1024p_feat --netG local --ngf 32 --num_D 3 --load_pretrain checkpoints/label2city_512p_feat/ --niter_fix_global 20 --resize_or_crop crop --fineSize 896 --instance_feat --load_features
\ No newline at end of file
diff --git a/scripts/train_1024p_feat_24G.sh b/scripts/train_1024p_feat_24G.sh
new file mode 100755
index 0000000..399d720
--- /dev/null
+++ b/scripts/train_1024p_feat_24G.sh
@@ -0,0 +1,6 @@
+############## To train images at 2048 x 1024 resolution after training 1024 x 512 resolution models #############

+######## Using GPUs with 24G memory

+# First precompute feature maps and save them

+python precompute_feature_maps.py --name label2city_512p_feat;

+# Adding instances and encoded features

+python train.py --name label2city_1024p_feat --netG local --ngf 32 --num_D 3 --load_pretrain checkpoints/label2city_512p_feat/ --niter 50 --niter_decay 50 --niter_fix_global 10 --resize_or_crop none --instance_feat --load_features
\ No newline at end of file
diff --git a/scripts/train_512p.sh b/scripts/train_512p.sh
new file mode 100755
index 0000000..222c348
--- /dev/null
+++ b/scripts/train_512p.sh
@@ -0,0 +1,2 @@
+### Using labels only

+python train.py --name label2city_512p
\ No newline at end of file
diff --git a/scripts/train_512p_feat.sh b/scripts/train_512p_feat.sh
new file mode 100755
index 0000000..9d4859c
--- /dev/null
+++ b/scripts/train_512p_feat.sh
@@ -0,0 +1,2 @@
+### Adding instances and encoded features

+python train.py --name label2city_512p_feat --instance_feat
\ No newline at end of file
diff --git a/scripts/train_512p_multigpu.sh b/scripts/train_512p_multigpu.sh
new file mode 100755
index 0000000..16f0a1a
--- /dev/null
+++ b/scripts/train_512p_multigpu.sh
@@ -0,0 +1,2 @@
+######## Multi-GPU training example #######

+python train.py --name label2city_512p --batchSize 8 --gpu_ids 0,1,2,3,4,5,6,7
\ No newline at end of file
diff --git a/test.py b/test.py
new file mode 100755
index 0000000..d96ac10
--- /dev/null
+++ b/test.py
@@ -0,0 +1,37 @@
+### Copyright (C) 2017 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 time
+import os
+from collections import OrderedDict
+from options.test_options import TestOptions
+from data.data_loader import CreateDataLoader
+from models.models import create_model
+import util.util as util
+from util.visualizer import Visualizer
+from util import html
+
+opt = TestOptions().parse(save=False)
+opt.nThreads = 1   # test code only supports nThreads = 1
+opt.batchSize = 1  # test code only supports batchSize = 1
+opt.serial_batches = True  # no shuffle
+opt.no_flip = True  # no flip
+
+data_loader = CreateDataLoader(opt)
+dataset = data_loader.load_data()
+model = create_model(opt)
+visualizer = Visualizer(opt)
+# create website
+web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch))
+webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch))
+# test
+for i, data in enumerate(dataset):
+    if i >= opt.how_many:
+        break
+    generated = model.inference(data['label'], data['inst'])
+    visuals = OrderedDict([('input_label', util.tensor2label(data['label'][0], opt.label_nc)),
+                           ('synthesized_image', util.tensor2im(generated.data[0]))])
+    img_path = data['path']
+    print('process image... %s' % img_path)
+    visualizer.save_images(webpage, visuals, img_path)
+
+webpage.save()
diff --git a/train.py b/train.py
new file mode 100755
index 0000000..4965481
--- /dev/null
+++ b/train.py
@@ -0,0 +1,118 @@
+### Copyright (C) 2017 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 time
+from collections import OrderedDict
+from options.train_options import TrainOptions
+from data.data_loader import CreateDataLoader
+from models.models import create_model
+import util.util as util
+from util.visualizer import Visualizer
+import os
+import numpy as np
+import torch
+from torch.autograd import Variable
+
+opt = TrainOptions().parse()
+iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt')
+if opt.continue_train:
+    try:
+        start_epoch, epoch_iter = np.loadtxt(iter_path , delimiter=',', dtype=int)
+    except:
+        start_epoch, epoch_iter = 1, 0
+    print('Resuming from epoch %d at iteration %d' % (start_epoch, epoch_iter))        
+else:    
+    start_epoch, epoch_iter = 1, 0
+
+if opt.debug:
+    opt.display_freq = 1
+    opt.print_freq = 1
+    opt.niter = 1
+    opt.niter_decay = 0
+    opt.max_dataset_size = 10
+
+data_loader = CreateDataLoader(opt)
+dataset = data_loader.load_data()
+dataset_size = len(data_loader)
+print('#training images = %d' % dataset_size)
+
+model = create_model(opt)
+visualizer = Visualizer(opt)
+
+total_steps = (start_epoch-1) * dataset_size + epoch_iter    
+for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1):
+    epoch_start_time = time.time()
+    if epoch != start_epoch:
+        epoch_iter = epoch_iter % dataset_size
+    for i, data in enumerate(dataset, start=epoch_iter):
+        iter_start_time = time.time()
+        total_steps += opt.batchSize
+        epoch_iter += opt.batchSize
+
+        # whether to collect output images
+        save_fake = total_steps % opt.display_freq == 0
+
+        ############## Forward Pass ######################
+        losses, generated = model(Variable(data['label']), Variable(data['inst']), 
+            Variable(data['image']), Variable(data['feat']), infer=save_fake)
+
+        # sum per device losses
+        losses = [ torch.mean(x) if not isinstance(x, int) else x for x in losses ]
+        loss_dict = dict(zip(model.module.loss_names, losses))
+
+        # calculate final loss scalar
+        loss_D = (loss_dict['D_fake'] + loss_dict['D_real']) * 0.5
+        loss_G = loss_dict['G_GAN'] + loss_dict['G_GAN_Feat'] + loss_dict['G_VGG']
+
+        ############### Backward Pass ####################
+        # update generator weights
+        model.module.optimizer_G.zero_grad()
+        loss_G.backward()
+        model.module.optimizer_G.step()
+
+        # update discriminator weights
+        model.module.optimizer_D.zero_grad()
+        loss_D.backward()
+        model.module.optimizer_D.step()
+
+        #call(["nvidia-smi", "--format=csv", "--query-gpu=memory.used,memory.free"]) 
+
+        ############## Display results and errors ##########
+        ### print out errors
+        if total_steps % opt.print_freq == 0:
+            errors = {k: v.data[0] if not isinstance(v, (int,long,float)) else v for k, v in loss_dict.items()}
+            t = (time.time() - iter_start_time) / opt.batchSize
+            visualizer.print_current_errors(epoch, epoch_iter, errors, t)
+            visualizer.plot_current_errors(errors, total_steps)
+
+        ### display output images
+        if save_fake:
+            visuals = OrderedDict([('input_label', util.tensor2label(data['label'][0], opt.label_nc)),
+                                   ('synthesized_image', util.tensor2im(generated.data[0])),
+                                   ('real_image', util.tensor2im(data['image'][0]))])
+            visualizer.display_current_results(visuals, epoch, total_steps)
+
+        ### save latest model
+        if total_steps % opt.save_latest_freq == 0:
+            print('saving the latest model (epoch %d, total_steps %d)' % (epoch, total_steps))
+            model.module.save('latest')            
+            np.savetxt(iter_path, (epoch, epoch_iter), delimiter=',', fmt='%d')
+       
+    # end of epoch 
+	iter_end_time = time.time()
+    print('End of epoch %d / %d \t Time Taken: %d sec' %
+          (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time))
+
+    ### save model for this epoch
+    if epoch % opt.save_epoch_freq == 0:
+        print('saving the model at the end of epoch %d, iters %d' % (epoch, total_steps))        
+        model.module.save('latest')
+        model.module.save(epoch)
+        np.savetxt(iter_path, (epoch+1, 0), delimiter=',', fmt='%d')
+
+    ### instead of only training the local enhancer, train the entire network after certain iterations
+    if (opt.niter_fix_global != 0) and (epoch == opt.niter_fix_global):
+        model.module.update_fixed_params()
+
+    ### linearly decay learning rate after certain iterations
+    if epoch > opt.niter:
+        model.module.update_learning_rate()
diff --git a/util/__init__.py b/util/__init__.py
new file mode 100755
index 0000000..e69de29
--- /dev/null
+++ b/util/__init__.py
diff --git a/util/html.py b/util/html.py
new file mode 100755
index 0000000..a80aa59
--- /dev/null
+++ b/util/html.py
@@ -0,0 +1,63 @@
+import dominate
+from dominate.tags import *
+import os
+
+
+class HTML:
+    def __init__(self, web_dir, title, reflesh=0):
+        self.title = title
+        self.web_dir = web_dir
+        self.img_dir = os.path.join(self.web_dir, 'images')
+        if not os.path.exists(self.web_dir):
+            os.makedirs(self.web_dir)
+        if not os.path.exists(self.img_dir):
+            os.makedirs(self.img_dir)
+
+        self.doc = dominate.document(title=title)
+        if reflesh > 0:
+            with self.doc.head:
+                meta(http_equiv="reflesh", content=str(reflesh))
+
+    def get_image_dir(self):
+        return self.img_dir
+
+    def add_header(self, str):
+        with self.doc:
+            h3(str)
+
+    def add_table(self, border=1):
+        self.t = table(border=border, style="table-layout: fixed;")
+        self.doc.add(self.t)
+
+    def add_images(self, ims, txts, links, width=512):
+        self.add_table()
+        with self.t:
+            with tr():
+                for im, txt, link in zip(ims, txts, links):
+                    with td(style="word-wrap: break-word;", halign="center", valign="top"):
+                        with p():
+                            with a(href=os.path.join('images', link)):
+                                img(style="width:%dpx" % (width), src=os.path.join('images', im))
+                            br()
+                            p(txt)
+
+    def save(self):
+        html_file = '%s/index.html' % self.web_dir
+        f = open(html_file, 'wt')
+        f.write(self.doc.render())
+        f.close()
+
+
+if __name__ == '__main__':
+    html = HTML('web/', 'test_html')
+    html.add_header('hello world')
+
+    ims = []
+    txts = []
+    links = []
+    for n in range(4):
+        ims.append('image_%d.jpg' % n)
+        txts.append('text_%d' % n)
+        links.append('image_%d.jpg' % n)
+    html.add_images(ims, txts, links)
+    html.save()
diff --git a/util/image_pool.py b/util/image_pool.py
new file mode 100755
index 0000000..152ef5b
--- /dev/null
+++ b/util/image_pool.py
@@ -0,0 +1,32 @@
+import random
+import numpy as np
+import torch
+from torch.autograd import Variable
+class ImagePool():
+    def __init__(self, pool_size):
+        self.pool_size = pool_size
+        if self.pool_size > 0:
+            self.num_imgs = 0
+            self.images = []
+
+    def query(self, images):
+        if self.pool_size == 0:
+            return images
+        return_images = []
+        for image in images.data:
+            image = torch.unsqueeze(image, 0)
+            if self.num_imgs < self.pool_size:
+                self.num_imgs = self.num_imgs + 1
+                self.images.append(image)
+                return_images.append(image)
+            else:
+                p = random.uniform(0, 1)
+                if p > 0.5:
+                    random_id = random.randint(0, self.pool_size-1)
+                    tmp = self.images[random_id].clone()
+                    self.images[random_id] = image
+                    return_images.append(tmp)
+                else:
+                    return_images.append(image)
+        return_images = Variable(torch.cat(return_images, 0))
+        return return_images
diff --git a/util/util.py b/util/util.py
new file mode 100755
index 0000000..0898f7a
--- /dev/null
+++ b/util/util.py
@@ -0,0 +1,99 @@
+### Copyright (C) 2017 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 __future__ import print_function
+import torch
+import numpy as np
+from PIL import Image
+import inspect, re
+import numpy as np
+import os
+import collections
+from PIL import Image
+
+# 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):
+    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
+    image_numpy = image_tensor.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)
+
+def tensor2label(output, n_label, imtype=np.uint8):
+    output = output.cpu().float()    
+    if output.size()[0] > 1:
+        output = output.max(0, keepdim=True)[1]
+    output = Colorize(n_label)(output)
+    output = np.transpose(output.numpy(), (1, 2, 0))
+    return output.astype(imtype)
+
+def save_image(image_numpy, image_path):
+    image_pil = Image.fromarray(image_numpy)
+    image_pil.save(image_path)
+
+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 uint82bin(n, count=8):
+    """returns the binary of integer n, count refers to amount of bits"""
+    return ''.join([str((n >> y) & 1) for y in range(count-1, -1, -1)])
+
+def labelcolormap(N):
+    if N == 35: # cityscape
+        cmap = np.array([(  0,  0,  0), (  0,  0,  0), (  0,  0,  0), (  0,  0,  0), (  0,  0,  0), (111, 74,  0), ( 81,  0, 81),
+                     (128, 64,128), (244, 35,232), (250,170,160), (230,150,140), ( 70, 70, 70), (102,102,156), (190,153,153),
+                     (180,165,180), (150,100,100), (150,120, 90), (153,153,153), (153,153,153), (250,170, 30), (220,220,  0),
+                     (107,142, 35), (152,251,152), ( 70,130,180), (220, 20, 60), (255,  0,  0), (  0,  0,142), (  0,  0, 70),
+                     (  0, 60,100), (  0,  0, 90), (  0,  0,110), (  0, 80,100), (  0,  0,230), (119, 11, 32), (  0,  0,142)], 
+                     dtype=np.uint8)
+    else:
+        cmap = np.zeros((N, 3), dtype=np.uint8)
+        for i in range(N):
+            r = 0
+            g = 0
+            b = 0
+            id = i
+            for j in range(7):
+                str_id = uint82bin(id)
+                r = r ^ (np.uint8(str_id[-1]) << (7-j))
+                g = g ^ (np.uint8(str_id[-2]) << (7-j))
+                b = b ^ (np.uint8(str_id[-3]) << (7-j))
+                id = id >> 3
+            cmap[i, 0] = r
+            cmap[i, 1] = g
+            cmap[i, 2] = b
+    return cmap
+
+class Colorize(object):
+    def __init__(self, n=35):
+        self.cmap = labelcolormap(n)
+        self.cmap = torch.from_numpy(self.cmap[:n])
+
+    def __call__(self, gray_image):
+        size = gray_image.size()
+        color_image = torch.ByteTensor(3, size[1], size[2]).fill_(0)
+
+        for label in range(0, len(self.cmap)):
+            mask = (label == gray_image[0]).cpu()
+            color_image[0][mask] = self.cmap[label][0]
+            color_image[1][mask] = self.cmap[label][1]
+            color_image[2][mask] = self.cmap[label][2]
+
+        return color_image
diff --git a/util/visualizer.py b/util/visualizer.py
new file mode 100755
index 0000000..f41c55a
--- /dev/null
+++ b/util/visualizer.py
@@ -0,0 +1,133 @@
+### Copyright (C) 2017 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
+import os
+import ntpath
+import time
+from . import util
+from . import html
+import scipy.misc 
+try:
+    from StringIO import StringIO  # Python 2.7
+except ImportError:
+    from io import BytesIO         # Python 3.x
+
+class Visualizer():
+    def __init__(self, opt):
+        # self.opt = opt
+        self.tf_log = opt.tf_log
+        self.use_html = opt.isTrain and not opt.no_html
+        self.win_size = opt.display_winsize
+        self.name = opt.name
+        if self.tf_log:
+            import tensorflow as tf
+            self.tf = tf
+            self.log_dir = os.path.join(opt.checkpoints_dir, opt.name, 'logs')
+            self.writer = tf.summary.FileWriter(self.log_dir)
+
+        if self.use_html:
+            self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web')
+            self.img_dir = os.path.join(self.web_dir, 'images')
+            print('create web directory %s...' % self.web_dir)
+            util.mkdirs([self.web_dir, self.img_dir])
+        self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')
+        with open(self.log_name, "a") as log_file:
+            now = time.strftime("%c")
+            log_file.write('================ Training Loss (%s) ================\n' % now)
+
+    # |visuals|: dictionary of images to display or save
+    def display_current_results(self, visuals, epoch, step):
+        if self.tf_log: # show images in tensorboard output
+            img_summaries = []
+            for label, image_numpy in visuals.items():
+                # Write the image to a string
+                try:
+                    s = StringIO()
+                except:
+                    s = BytesIO()
+                scipy.misc.toimage(image_numpy).save(s, format="jpeg")
+                # Create an Image object
+                img_sum = self.tf.Summary.Image(encoded_image_string=s.getvalue(), height=image_numpy.shape[0], width=image_numpy.shape[1])
+                # Create a Summary value
+                img_summaries.append(self.tf.Summary.Value(tag=label, image=img_sum))
+
+            # Create and write Summary
+            summary = self.tf.Summary(value=img_summaries)
+            self.writer.add_summary(summary, step)
+
+        if self.use_html: # save images to a html file
+            for label, image_numpy in visuals.items():
+                if isinstance(image_numpy, list):
+                    for i in range(len(image_numpy)):
+                        img_path = os.path.join(self.img_dir, 'epoch%.3d_%s_%d.jpg' % (epoch, label, i))
+                        util.save_image(image_numpy[i], img_path)
+                else:
+                    img_path = os.path.join(self.img_dir, 'epoch%.3d_%s.jpg' % (epoch, label))
+                    util.save_image(image_numpy, img_path)
+
+            # update website
+            webpage = html.HTML(self.web_dir, 'Experiment name = %s' % self.name, reflesh=1)
+            for n in range(epoch, 0, -1):
+                webpage.add_header('epoch [%d]' % n)
+                ims = []
+                txts = []
+                links = []
+
+                for label, image_numpy in visuals.items():
+                    if isinstance(image_numpy, list):
+                        for i in range(len(image_numpy)):
+                            img_path = 'epoch%.3d_%s_%d.jpg' % (n, label, i)
+                            ims.append(img_path)
+                            txts.append(label+str(i))
+                            links.append(img_path)
+                    else:
+                        img_path = 'epoch%.3d_%s.jpg' % (n, label)
+                        ims.append(img_path)
+                        txts.append(label)
+                        links.append(img_path)
+                if len(ims) < 10:
+                    webpage.add_images(ims, txts, links, width=self.win_size)
+                else:
+                    num = int(round(len(ims)/2.0))
+                    webpage.add_images(ims[:num], txts[:num], links[:num], width=self.win_size)
+                    webpage.add_images(ims[num:], txts[num:], links[num:], width=self.win_size)
+            webpage.save()
+
+    # errors: dictionary of error labels and values
+    def plot_current_errors(self, errors, step):
+        if self.tf_log:
+            for tag, value in errors.items():            
+                summary = self.tf.Summary(value=[self.tf.Summary.Value(tag=tag, simple_value=value)])
+                self.writer.add_summary(summary, step)
+
+    # errors: same format as |errors| of plotCurrentErrors
+    def print_current_errors(self, epoch, i, errors, t):
+        message = '(epoch: %d, iters: %d, time: %.3f) ' % (epoch, i, t)
+        for k, v in errors.items():
+            if v != 0:
+                message += '%s: %.3f ' % (k, v)
+
+        print(message)
+        with open(self.log_name, "a") as log_file:
+            log_file.write('%s\n' % message)
+
+    # save image to the disk
+    def save_images(self, webpage, visuals, image_path):
+        image_dir = webpage.get_image_dir()
+        short_path = ntpath.basename(image_path[0])
+        name = os.path.splitext(short_path)[0]
+
+        webpage.add_header(name)
+        ims = []
+        txts = []
+        links = []
+
+        for label, image_numpy in visuals.items():
+            image_name = '%s_%s.jpg' % (name, label)
+            save_path = os.path.join(image_dir, image_name)
+            util.save_image(image_numpy, save_path)
+
+            ims.append(image_name)
+            txts.append(label)
+            links.append(image_name)
+        webpage.add_images(ims, txts, links, width=self.win_size)