Initial commit

scripts/train_1024p_12G.sh

+############## 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

scripts/train_1024p_24G.sh

+############## 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

scripts/train_1024p_feat_12G.sh

+############## 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

scripts/train_1024p_feat_24G.sh

+############## 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

scripts/train_512p.sh

+### Using labels only
+python train.py --name label2city_512p
\ No newline at end of file

scripts/train_512p_feat.sh

+### Adding instances and encoded features
+python train.py --name label2city_512p_feat --instance_feat
\ No newline at end of file

scripts/train_512p_fp16.sh

+### Using labels only
+ python -m torch.distributed.launch train.py --name label2city_512p --fp16
\ No newline at end of file

scripts/train_512p_fp16_multigpu.sh

+######## Multi-GPU training example #######
+python -m torch.distributed.launch train.py --name label2city_512p --batchSize 8 --gpu_ids 0,1,2,3,4,5,6,7 --fp16
\ No newline at end of file

scripts/train_512p_multigpu.sh

+######## 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

test.py

+import os
+from collections import OrderedDict
+from torch.autograd import Variable
+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
+import torch
+
+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()
+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
+if not opt.engine and not opt.onnx:
+    model = create_model(opt)
+    if opt.data_type == 16:
+        model.half()
+    elif opt.data_type == 8:
+        model.type(torch.uint8)
+            
+    if opt.verbose:
+        print(model)
+else:
+    from run_engine import run_trt_engine, run_onnx
+    
+for i, data in enumerate(dataset):
+    if i >= opt.how_many:
+        break
+    if opt.data_type == 16:
+        data['label'] = data['label'].half()
+        data['inst']  = data['inst'].half()
+    elif opt.data_type == 8:
+        data['label'] = data['label'].uint8()
+        data['inst']  = data['inst'].uint8()
+    if opt.export_onnx:
+        print ("Exporting to ONNX: ", opt.export_onnx)
+        assert opt.export_onnx.endswith("onnx"), "Export model file should end with .onnx"
+        torch.onnx.export(model, [data['label'], data['inst']],
+                          opt.export_onnx, verbose=True)
+        exit(0)
+    minibatch = 1 
+    if opt.engine:
+        generated = run_trt_engine(opt.engine, minibatch, [data['label'], data['inst']])
+    elif opt.onnx:
+        generated = run_onnx(opt.onnx, opt.data_type, minibatch, [data['label'], data['inst']])
+    else:        
+        generated = model.inference(data['label'], data['inst'], data['image'])
+        
+    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()

train.py

+import time
+import os
+import numpy as np
+import torch
+from torch.autograd import Variable
+from collections import OrderedDict
+from subprocess import call
+import fractions
+def lcm(a,b): return abs(a * b)/fractions.gcd(a,b) if a and b else 0
+
+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
+
+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
+
+opt.print_freq = lcm(opt.print_freq, opt.batchSize)    
+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)
+if opt.fp16:    
+    from apex import amp
+    model, [optimizer_G, optimizer_D] = amp.initialize(model, [model.optimizer_G, model.optimizer_D], opt_level='O1')             
+    model = torch.nn.DataParallel(model, device_ids=opt.gpu_ids)
+else:
+    optimizer_G, optimizer_D = model.module.optimizer_G, model.module.optimizer_D
+
+total_steps = (start_epoch-1) * dataset_size + epoch_iter
+
+display_delta = total_steps % opt.display_freq
+print_delta = total_steps % opt.print_freq
+save_delta = total_steps % opt.save_latest_freq
+
+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):
+        if total_steps % opt.print_freq == print_delta:
+            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 == display_delta
+
+        ############## 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.get('G_GAN_Feat',0) + loss_dict.get('G_VGG',0)
+
+        ############### Backward Pass ####################
+        # update generator weights
+        optimizer_G.zero_grad()
+        if opt.fp16:                                
+            with amp.scale_loss(loss_G, optimizer_G) as scaled_loss: scaled_loss.backward()                
+        else:
+            loss_G.backward()          
+        optimizer_G.step()
+
+        # update discriminator weights
+        optimizer_D.zero_grad()
+        if opt.fp16:                                
+            with amp.scale_loss(loss_D, optimizer_D) as scaled_loss: scaled_loss.backward()                
+        else:
+            loss_D.backward()        
+        optimizer_D.step()        
+
+        ############## Display results and errors ##########
+        ### print out errors
+        if total_steps % opt.print_freq == print_delta:
+            errors = {k: v.data.item() if not isinstance(v, int) else v for k, v in loss_dict.items()}            
+            t = (time.time() - iter_start_time) / opt.print_freq
+            visualizer.print_current_errors(epoch, epoch_iter, errors, t)
+            visualizer.plot_current_errors(errors, total_steps)
+            #call(["nvidia-smi", "--format=csv", "--query-gpu=memory.used,memory.free"]) 
+
+        ### 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 == save_delta:
+            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')
+
+        if epoch_iter >= dataset_size:
+            break
+       
+    # 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()

util/html.py

+import dominate
+from dominate.tags import *
+import os
+
+
+class HTML:
+    def __init__(self, web_dir, title, refresh=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 refresh > 0:
+            with self.doc.head:
+                meta(http_equiv="refresh", content=str(refresh))
+
+    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()

util/image_pool.py

+import random
+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

util/util.py

+from __future__ import print_function
+import torch
+import numpy as np
+from PIL import Image
+import numpy as np
+import os
+
+# Converts a Tensor into a Numpy array
+# |imtype|: the desired type of the converted numpy array
+def tensor2im(image_tensor, imtype=np.uint8, normalize=True):
+    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 or image_numpy.shape[2] > 3:        
+        image_numpy = image_numpy[:,:,0]
+    return image_numpy.astype(imtype)
+
+# Converts a one-hot tensor into a colorful label map
+def tensor2label(label_tensor, n_label, imtype=np.uint8):
+    if n_label == 0:
+        return tensor2im(label_tensor, imtype)
+    label_tensor = label_tensor.cpu().float()    
+    if label_tensor.size()[0] > 1:
+        label_tensor = label_tensor.max(0, keepdim=True)[1]
+    label_tensor = Colorize(n_label)(label_tensor)
+    label_numpy = np.transpose(label_tensor.numpy(), (1, 2, 0))
+    return label_numpy.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)
+
+###############################################################################
+# Code from
+# https://github.com/ycszen/pytorch-seg/blob/master/transform.py
+# Modified so it complies with the Citscape label map colors
+###############################################################################
+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, g, b = 0, 0, 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

util/visualizer.py

+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, refresh=30)
+            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)