\n\n![](\"images/teaser_high_res_uncompressed.png\")
\n\n\n\n\n\t
\n\n\n\n\n![\"paper](\"images/paper_thumbnail.jpg\")
\n\n\n\n
\n\n\n\t
\n\n
\n![](\"images/method.png\")
\n
\n![](\"images/coco_comparison.jpg\")
\n\n
\n![](\"images/flickr.jpg\")
\n\n
\n
\n\n
\n 0:\n assert(torch.cuda.is_available())\n net.cuda()\n net.init_weights(opt.init_type, opt.init_variance)\n return net\n\n\ndef define_G(opt):\n netG_cls = find_network_using_name(opt.netG, 'generator')\n return create_network(netG_cls, opt)\n\n\ndef define_D(opt):\n netD_cls = find_network_using_name(opt.netD, 'discriminator')\n return create_network(netD_cls, opt)\n\n\ndef define_E(opt):\n # there exists only one encoder type\n netE_cls = find_network_using_name('conv', 'encoder')\n return create_network(netE_cls, opt)\n"
},
{
"path": "models/networks/architecture.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nimport torch.nn.utils.spectral_norm as spectral_norm\nfrom models.networks.normalization import SPADE\n\n\n# ResNet block that uses SPADE.\n# It differs from the ResNet block of pix2pixHD in that\n# it takes in the segmentation map as input, learns the skip connection if necessary,\n# and applies normalization first and then convolution.\n# This architecture seemed like a standard architecture for unconditional or\n# class-conditional GAN architecture using residual block.\n# The code was inspired from https://github.com/LMescheder/GAN_stability.\nclass SPADEResnetBlock(nn.Module):\n def __init__(self, fin, fout, opt):\n super().__init__()\n # Attributes\n self.learned_shortcut = (fin != fout)\n fmiddle = min(fin, fout)\n\n # create conv layers\n self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=1)\n self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=1)\n if self.learned_shortcut:\n self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False)\n\n # apply spectral norm if specified\n if 'spectral' in opt.norm_G:\n self.conv_0 = spectral_norm(self.conv_0)\n self.conv_1 = spectral_norm(self.conv_1)\n if self.learned_shortcut:\n self.conv_s = spectral_norm(self.conv_s)\n\n # define normalization layers\n spade_config_str = opt.norm_G.replace('spectral', '')\n self.norm_0 = SPADE(spade_config_str, fin, opt.semantic_nc)\n self.norm_1 = SPADE(spade_config_str, fmiddle, opt.semantic_nc)\n if self.learned_shortcut:\n self.norm_s = SPADE(spade_config_str, fin, opt.semantic_nc)\n\n # note the resnet block with SPADE also takes in |seg|,\n # the semantic segmentation map as input\n def forward(self, x, seg):\n x_s = self.shortcut(x, seg)\n\n dx = self.conv_0(self.actvn(self.norm_0(x, seg)))\n dx = self.conv_1(self.actvn(self.norm_1(dx, seg)))\n\n out = x_s + dx\n\n return out\n\n def shortcut(self, x, seg):\n if self.learned_shortcut:\n x_s = self.conv_s(self.norm_s(x, seg))\n else:\n x_s = x\n return x_s\n\n def actvn(self, x):\n return F.leaky_relu(x, 2e-1)\n\n\n# ResNet block used in pix2pixHD\n# We keep the same architecture as pix2pixHD.\nclass ResnetBlock(nn.Module):\n def __init__(self, dim, norm_layer, activation=nn.ReLU(False), kernel_size=3):\n super().__init__()\n\n pw = (kernel_size - 1) // 2\n self.conv_block = nn.Sequential(\n nn.ReflectionPad2d(pw),\n norm_layer(nn.Conv2d(dim, dim, kernel_size=kernel_size)),\n activation,\n nn.ReflectionPad2d(pw),\n norm_layer(nn.Conv2d(dim, dim, kernel_size=kernel_size))\n )\n\n def forward(self, x):\n y = self.conv_block(x)\n out = x + y\n return out\n\n\n# VGG architecter, used for the perceptual loss using a pretrained VGG network\nclass VGG19(torch.nn.Module):\n def __init__(self, requires_grad=False):\n super().__init__()\n vgg_pretrained_features = torchvision.models.vgg19(pretrained=True).features\n self.slice1 = torch.nn.Sequential()\n self.slice2 = torch.nn.Sequential()\n self.slice3 = torch.nn.Sequential()\n self.slice4 = torch.nn.Sequential()\n self.slice5 = torch.nn.Sequential()\n for x in range(2):\n self.slice1.add_module(str(x), vgg_pretrained_features[x])\n for x in range(2, 7):\n self.slice2.add_module(str(x), vgg_pretrained_features[x])\n for x in range(7, 12):\n self.slice3.add_module(str(x), vgg_pretrained_features[x])\n for x in range(12, 21):\n self.slice4.add_module(str(x), vgg_pretrained_features[x])\n for x in range(21, 30):\n self.slice5.add_module(str(x), vgg_pretrained_features[x])\n if not requires_grad:\n for param in self.parameters():\n param.requires_grad = False\n\n def forward(self, X):\n h_relu1 = self.slice1(X)\n h_relu2 = self.slice2(h_relu1)\n h_relu3 = self.slice3(h_relu2)\n h_relu4 = self.slice4(h_relu3)\n h_relu5 = self.slice5(h_relu4)\n out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]\n return out\n"
},
{
"path": "models/networks/base_network.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nimport torch.nn as nn\nfrom torch.nn import init\n\n\nclass BaseNetwork(nn.Module):\n def __init__(self):\n super(BaseNetwork, self).__init__()\n\n @staticmethod\n def modify_commandline_options(parser, is_train):\n return parser\n\n def print_network(self):\n if isinstance(self, list):\n self = self[0]\n num_params = 0\n for param in self.parameters():\n num_params += param.numel()\n print('Network [%s] was created. Total number of parameters: %.1f million. '\n 'To see the architecture, do print(network).'\n % (type(self).__name__, num_params / 1000000))\n\n def init_weights(self, init_type='normal', gain=0.02):\n def init_func(m):\n classname = m.__class__.__name__\n if classname.find('BatchNorm2d') != -1:\n if hasattr(m, 'weight') and m.weight is not None:\n init.normal_(m.weight.data, 1.0, gain)\n if hasattr(m, 'bias') and m.bias is not None:\n init.constant_(m.bias.data, 0.0)\n elif hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):\n if init_type == 'normal':\n init.normal_(m.weight.data, 0.0, gain)\n elif init_type == 'xavier':\n init.xavier_normal_(m.weight.data, gain=gain)\n elif init_type == 'xavier_uniform':\n init.xavier_uniform_(m.weight.data, gain=1.0)\n elif init_type == 'kaiming':\n init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')\n elif init_type == 'orthogonal':\n init.orthogonal_(m.weight.data, gain=gain)\n elif init_type == 'none': # uses pytorch's default init method\n m.reset_parameters()\n else:\n raise NotImplementedError('initialization method [%s] is not implemented' % init_type)\n if hasattr(m, 'bias') and m.bias is not None:\n init.constant_(m.bias.data, 0.0)\n\n self.apply(init_func)\n\n # propagate to children\n for m in self.children():\n if hasattr(m, 'init_weights'):\n m.init_weights(init_type, gain)\n"
},
{
"path": "models/networks/discriminator.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nimport torch.nn as nn\nimport numpy as np\nimport torch.nn.functional as F\nfrom models.networks.base_network import BaseNetwork\nfrom models.networks.normalization import get_nonspade_norm_layer\nimport util.util as util\n\n\nclass MultiscaleDiscriminator(BaseNetwork):\n @staticmethod\n def modify_commandline_options(parser, is_train):\n parser.add_argument('--netD_subarch', type=str, default='n_layer',\n help='architecture of each discriminator')\n parser.add_argument('--num_D', type=int, default=2,\n help='number of discriminators to be used in multiscale')\n opt, _ = parser.parse_known_args()\n\n # define properties of each discriminator of the multiscale discriminator\n subnetD = util.find_class_in_module(opt.netD_subarch + 'discriminator',\n 'models.networks.discriminator')\n subnetD.modify_commandline_options(parser, is_train)\n\n return parser\n\n def __init__(self, opt):\n super().__init__()\n self.opt = opt\n\n for i in range(opt.num_D):\n subnetD = self.create_single_discriminator(opt)\n self.add_module('discriminator_%d' % i, subnetD)\n\n def create_single_discriminator(self, opt):\n subarch = opt.netD_subarch\n if subarch == 'n_layer':\n netD = NLayerDiscriminator(opt)\n else:\n raise ValueError('unrecognized discriminator subarchitecture %s' % subarch)\n return netD\n\n def downsample(self, input):\n return F.avg_pool2d(input, kernel_size=3,\n stride=2, padding=[1, 1],\n count_include_pad=False)\n\n # Returns list of lists of discriminator outputs.\n # The final result is of size opt.num_D x opt.n_layers_D\n def forward(self, input):\n result = []\n get_intermediate_features = not self.opt.no_ganFeat_loss\n for name, D in self.named_children():\n out = D(input)\n if not get_intermediate_features:\n out = [out]\n result.append(out)\n input = self.downsample(input)\n\n return result\n\n\n# Defines the PatchGAN discriminator with the specified arguments.\nclass NLayerDiscriminator(BaseNetwork):\n @staticmethod\n def modify_commandline_options(parser, is_train):\n parser.add_argument('--n_layers_D', type=int, default=4,\n help='# layers in each discriminator')\n return parser\n\n def __init__(self, opt):\n super().__init__()\n self.opt = opt\n\n kw = 4\n padw = int(np.ceil((kw - 1.0) / 2))\n nf = opt.ndf\n input_nc = self.compute_D_input_nc(opt)\n\n norm_layer = get_nonspade_norm_layer(opt, opt.norm_D)\n sequence = [[nn.Conv2d(input_nc, nf, kernel_size=kw, stride=2, padding=padw),\n nn.LeakyReLU(0.2, False)]]\n\n for n in range(1, opt.n_layers_D):\n nf_prev = nf\n nf = min(nf * 2, 512)\n stride = 1 if n == opt.n_layers_D - 1 else 2\n sequence += [[norm_layer(nn.Conv2d(nf_prev, nf, kernel_size=kw,\n stride=stride, padding=padw)),\n nn.LeakyReLU(0.2, False)\n ]]\n\n sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)]]\n\n # We divide the layers into groups to extract intermediate layer outputs\n for n in range(len(sequence)):\n self.add_module('model' + str(n), nn.Sequential(*sequence[n]))\n\n def compute_D_input_nc(self, opt):\n input_nc = opt.label_nc + opt.output_nc\n if opt.contain_dontcare_label:\n input_nc += 1\n if not opt.no_instance:\n input_nc += 1\n return input_nc\n\n def forward(self, input):\n results = [input]\n for submodel in self.children():\n intermediate_output = submodel(results[-1])\n results.append(intermediate_output)\n\n get_intermediate_features = not self.opt.no_ganFeat_loss\n if get_intermediate_features:\n return results[1:]\n else:\n return results[-1]\n"
},
{
"path": "models/networks/encoder.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nimport torch.nn as nn\nimport numpy as np\nimport torch.nn.functional as F\nfrom models.networks.base_network import BaseNetwork\nfrom models.networks.normalization import get_nonspade_norm_layer\n\n\nclass ConvEncoder(BaseNetwork):\n \"\"\" Same architecture as the image discriminator \"\"\"\n\n def __init__(self, opt):\n super().__init__()\n\n kw = 3\n pw = int(np.ceil((kw - 1.0) / 2))\n ndf = opt.ngf\n norm_layer = get_nonspade_norm_layer(opt, opt.norm_E)\n self.layer1 = norm_layer(nn.Conv2d(3, ndf, kw, stride=2, padding=pw))\n self.layer2 = norm_layer(nn.Conv2d(ndf * 1, ndf * 2, kw, stride=2, padding=pw))\n self.layer3 = norm_layer(nn.Conv2d(ndf * 2, ndf * 4, kw, stride=2, padding=pw))\n self.layer4 = norm_layer(nn.Conv2d(ndf * 4, ndf * 8, kw, stride=2, padding=pw))\n self.layer5 = norm_layer(nn.Conv2d(ndf * 8, ndf * 8, kw, stride=2, padding=pw))\n if opt.crop_size >= 256:\n self.layer6 = norm_layer(nn.Conv2d(ndf * 8, ndf * 8, kw, stride=2, padding=pw))\n\n self.so = s0 = 4\n self.fc_mu = nn.Linear(ndf * 8 * s0 * s0, 256)\n self.fc_var = nn.Linear(ndf * 8 * s0 * s0, 256)\n\n self.actvn = nn.LeakyReLU(0.2, False)\n self.opt = opt\n\n def forward(self, x):\n if x.size(2) != 256 or x.size(3) != 256:\n x = F.interpolate(x, size=(256, 256), mode='bilinear')\n\n x = self.layer1(x)\n x = self.layer2(self.actvn(x))\n x = self.layer3(self.actvn(x))\n x = self.layer4(self.actvn(x))\n x = self.layer5(self.actvn(x))\n if self.opt.crop_size >= 256:\n x = self.layer6(self.actvn(x))\n x = self.actvn(x)\n\n x = x.view(x.size(0), -1)\n mu = self.fc_mu(x)\n logvar = self.fc_var(x)\n\n return mu, logvar\n"
},
{
"path": "models/networks/generator.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom models.networks.base_network import BaseNetwork\nfrom models.networks.normalization import get_nonspade_norm_layer\nfrom models.networks.architecture import ResnetBlock as ResnetBlock\nfrom models.networks.architecture import SPADEResnetBlock as SPADEResnetBlock\n\n\nclass SPADEGenerator(BaseNetwork):\n @staticmethod\n def modify_commandline_options(parser, is_train):\n parser.set_defaults(norm_G='spectralspadesyncbatch3x3')\n parser.add_argument('--num_upsampling_layers',\n choices=('normal', 'more', 'most'), default='normal',\n help=\"If 'more', adds upsampling layer between the two middle resnet blocks. If 'most', also add one more upsampling + resnet layer at the end of the generator\")\n\n return parser\n\n def __init__(self, opt):\n super().__init__()\n self.opt = opt\n nf = opt.ngf\n\n self.sw, self.sh = self.compute_latent_vector_size(opt)\n\n if opt.use_vae:\n # In case of VAE, we will sample from random z vector\n self.fc = nn.Linear(opt.z_dim, 16 * nf * self.sw * self.sh)\n else:\n # Otherwise, we make the network deterministic by starting with\n # downsampled segmentation map instead of random z\n self.fc = nn.Conv2d(self.opt.semantic_nc, 16 * nf, 3, padding=1)\n\n self.head_0 = SPADEResnetBlock(16 * nf, 16 * nf, opt)\n\n self.G_middle_0 = SPADEResnetBlock(16 * nf, 16 * nf, opt)\n self.G_middle_1 = SPADEResnetBlock(16 * nf, 16 * nf, opt)\n\n self.up_0 = SPADEResnetBlock(16 * nf, 8 * nf, opt)\n self.up_1 = SPADEResnetBlock(8 * nf, 4 * nf, opt)\n self.up_2 = SPADEResnetBlock(4 * nf, 2 * nf, opt)\n self.up_3 = SPADEResnetBlock(2 * nf, 1 * nf, opt)\n\n final_nc = nf\n\n if opt.num_upsampling_layers == 'most':\n self.up_4 = SPADEResnetBlock(1 * nf, nf // 2, opt)\n final_nc = nf // 2\n\n self.conv_img = nn.Conv2d(final_nc, 3, 3, padding=1)\n\n self.up = nn.Upsample(scale_factor=2)\n\n def compute_latent_vector_size(self, opt):\n if opt.num_upsampling_layers == 'normal':\n num_up_layers = 5\n elif opt.num_upsampling_layers == 'more':\n num_up_layers = 6\n elif opt.num_upsampling_layers == 'most':\n num_up_layers = 7\n else:\n raise ValueError('opt.num_upsampling_layers [%s] not recognized' %\n opt.num_upsampling_layers)\n\n sw = opt.crop_size // (2**num_up_layers)\n sh = round(sw / opt.aspect_ratio)\n\n return sw, sh\n\n def forward(self, input, z=None):\n seg = input\n\n if self.opt.use_vae:\n # we sample z from unit normal and reshape the tensor\n if z is None:\n z = torch.randn(input.size(0), self.opt.z_dim,\n dtype=torch.float32, device=input.get_device())\n x = self.fc(z)\n x = x.view(-1, 16 * self.opt.ngf, self.sh, self.sw)\n else:\n # we downsample segmap and run convolution\n x = F.interpolate(seg, size=(self.sh, self.sw))\n x = self.fc(x)\n\n x = self.head_0(x, seg)\n\n x = self.up(x)\n x = self.G_middle_0(x, seg)\n\n if self.opt.num_upsampling_layers == 'more' or \\\n self.opt.num_upsampling_layers == 'most':\n x = self.up(x)\n\n x = self.G_middle_1(x, seg)\n\n x = self.up(x)\n x = self.up_0(x, seg)\n x = self.up(x)\n x = self.up_1(x, seg)\n x = self.up(x)\n x = self.up_2(x, seg)\n x = self.up(x)\n x = self.up_3(x, seg)\n\n if self.opt.num_upsampling_layers == 'most':\n x = self.up(x)\n x = self.up_4(x, seg)\n\n x = self.conv_img(F.leaky_relu(x, 2e-1))\n x = F.tanh(x)\n\n return x\n\n\nclass Pix2PixHDGenerator(BaseNetwork):\n @staticmethod\n def modify_commandline_options(parser, is_train):\n parser.add_argument('--resnet_n_downsample', type=int, default=4, help='number of downsampling layers in netG')\n parser.add_argument('--resnet_n_blocks', type=int, default=9, help='number of residual blocks in the global generator network')\n parser.add_argument('--resnet_kernel_size', type=int, default=3,\n help='kernel size of the resnet block')\n parser.add_argument('--resnet_initial_kernel_size', type=int, default=7,\n help='kernel size of the first convolution')\n parser.set_defaults(norm_G='instance')\n return parser\n\n def __init__(self, opt):\n super().__init__()\n input_nc = opt.label_nc + (1 if opt.contain_dontcare_label else 0) + (0 if opt.no_instance else 1)\n\n norm_layer = get_nonspade_norm_layer(opt, opt.norm_G)\n activation = nn.ReLU(False)\n\n model = []\n\n # initial conv\n model += [nn.ReflectionPad2d(opt.resnet_initial_kernel_size // 2),\n norm_layer(nn.Conv2d(input_nc, opt.ngf,\n kernel_size=opt.resnet_initial_kernel_size,\n padding=0)),\n activation]\n\n # downsample\n mult = 1\n for i in range(opt.resnet_n_downsample):\n model += [norm_layer(nn.Conv2d(opt.ngf * mult, opt.ngf * mult * 2,\n kernel_size=3, stride=2, padding=1)),\n activation]\n mult *= 2\n\n # resnet blocks\n for i in range(opt.resnet_n_blocks):\n model += [ResnetBlock(opt.ngf * mult,\n norm_layer=norm_layer,\n activation=activation,\n kernel_size=opt.resnet_kernel_size)]\n\n # upsample\n for i in range(opt.resnet_n_downsample):\n nc_in = int(opt.ngf * mult)\n nc_out = int((opt.ngf * mult) / 2)\n model += [norm_layer(nn.ConvTranspose2d(nc_in, nc_out,\n kernel_size=3, stride=2,\n padding=1, output_padding=1)),\n activation]\n mult = mult // 2\n\n # final output conv\n model += [nn.ReflectionPad2d(3),\n nn.Conv2d(nc_out, opt.output_nc, kernel_size=7, padding=0),\n nn.Tanh()]\n\n self.model = nn.Sequential(*model)\n\n def forward(self, input, z=None):\n return self.model(input)\n"
},
{
"path": "models/networks/loss.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom models.networks.architecture import VGG19\n\n\n# Defines the GAN loss which uses either LSGAN or the regular GAN.\n# When LSGAN is used, it is basically same as MSELoss,\n# but it abstracts away the need to create the target label tensor\n# that has the same size as the input\nclass GANLoss(nn.Module):\n def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0,\n tensor=torch.FloatTensor, opt=None):\n super(GANLoss, self).__init__()\n self.real_label = target_real_label\n self.fake_label = target_fake_label\n self.real_label_tensor = None\n self.fake_label_tensor = None\n self.zero_tensor = None\n self.Tensor = tensor\n self.gan_mode = gan_mode\n self.opt = opt\n if gan_mode == 'ls':\n pass\n elif gan_mode == 'original':\n pass\n elif gan_mode == 'w':\n pass\n elif gan_mode == 'hinge':\n pass\n else:\n raise ValueError('Unexpected gan_mode {}'.format(gan_mode))\n\n def get_target_tensor(self, input, target_is_real):\n if target_is_real:\n if self.real_label_tensor is None:\n self.real_label_tensor = self.Tensor(1).fill_(self.real_label)\n self.real_label_tensor.requires_grad_(False)\n return self.real_label_tensor.expand_as(input)\n else:\n if self.fake_label_tensor is None:\n self.fake_label_tensor = self.Tensor(1).fill_(self.fake_label)\n self.fake_label_tensor.requires_grad_(False)\n return self.fake_label_tensor.expand_as(input)\n\n def get_zero_tensor(self, input):\n if self.zero_tensor is None:\n self.zero_tensor = self.Tensor(1).fill_(0)\n self.zero_tensor.requires_grad_(False)\n return self.zero_tensor.expand_as(input)\n\n def loss(self, input, target_is_real, for_discriminator=True):\n if self.gan_mode == 'original': # cross entropy loss\n target_tensor = self.get_target_tensor(input, target_is_real)\n loss = F.binary_cross_entropy_with_logits(input, target_tensor)\n return loss\n elif self.gan_mode == 'ls':\n target_tensor = self.get_target_tensor(input, target_is_real)\n return F.mse_loss(input, target_tensor)\n elif self.gan_mode == 'hinge':\n if for_discriminator:\n if target_is_real:\n minval = torch.min(input - 1, self.get_zero_tensor(input))\n loss = -torch.mean(minval)\n else:\n minval = torch.min(-input - 1, self.get_zero_tensor(input))\n loss = -torch.mean(minval)\n else:\n assert target_is_real, \"The generator's hinge loss must be aiming for real\"\n loss = -torch.mean(input)\n return loss\n else:\n # wgan\n if target_is_real:\n return -input.mean()\n else:\n return input.mean()\n\n def __call__(self, input, target_is_real, for_discriminator=True):\n # computing loss is a bit complicated because |input| may not be\n # a tensor, but list of tensors in case of multiscale discriminator\n if isinstance(input, list):\n loss = 0\n for pred_i in input:\n if isinstance(pred_i, list):\n pred_i = pred_i[-1]\n loss_tensor = self.loss(pred_i, target_is_real, for_discriminator)\n bs = 1 if len(loss_tensor.size()) == 0 else loss_tensor.size(0)\n new_loss = torch.mean(loss_tensor.view(bs, -1), dim=1)\n loss += new_loss\n return loss / len(input)\n else:\n return self.loss(input, target_is_real, for_discriminator)\n\n\n# Perceptual loss that uses a pretrained VGG network\nclass VGGLoss(nn.Module):\n def __init__(self, gpu_ids):\n super(VGGLoss, self).__init__()\n self.vgg = VGG19().cuda()\n self.criterion = nn.L1Loss()\n self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]\n\n def forward(self, x, y):\n x_vgg, y_vgg = self.vgg(x), self.vgg(y)\n loss = 0\n for i in range(len(x_vgg)):\n loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())\n return loss\n\n\n# KL Divergence loss used in VAE with an image encoder\nclass KLDLoss(nn.Module):\n def forward(self, mu, logvar):\n return -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())\n"
},
{
"path": "models/networks/normalization.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nimport re\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom models.networks.sync_batchnorm import SynchronizedBatchNorm2d\nimport torch.nn.utils.spectral_norm as spectral_norm\n\n\n# Returns a function that creates a normalization function\n# that does not condition on semantic map\ndef get_nonspade_norm_layer(opt, norm_type='instance'):\n # helper function to get # output channels of the previous layer\n def get_out_channel(layer):\n if hasattr(layer, 'out_channels'):\n return getattr(layer, 'out_channels')\n return layer.weight.size(0)\n\n # this function will be returned\n def add_norm_layer(layer):\n nonlocal norm_type\n if norm_type.startswith('spectral'):\n layer = spectral_norm(layer)\n subnorm_type = norm_type[len('spectral'):]\n\n if subnorm_type == 'none' or len(subnorm_type) == 0:\n return layer\n\n # remove bias in the previous layer, which is meaningless\n # since it has no effect after normalization\n if getattr(layer, 'bias', None) is not None:\n delattr(layer, 'bias')\n layer.register_parameter('bias', None)\n\n if subnorm_type == 'batch':\n norm_layer = nn.BatchNorm2d(get_out_channel(layer), affine=True)\n elif subnorm_type == 'sync_batch':\n norm_layer = SynchronizedBatchNorm2d(get_out_channel(layer), affine=True)\n elif subnorm_type == 'instance':\n norm_layer = nn.InstanceNorm2d(get_out_channel(layer), affine=False)\n else:\n raise ValueError('normalization layer %s is not recognized' % subnorm_type)\n\n return nn.Sequential(layer, norm_layer)\n\n return add_norm_layer\n\n\n# Creates SPADE normalization layer based on the given configuration\n# SPADE consists of two steps. First, it normalizes the activations using\n# your favorite normalization method, such as Batch Norm or Instance Norm.\n# Second, it applies scale and bias to the normalized output, conditioned on\n# the segmentation map.\n# The format of |config_text| is spade(norm)(ks), where\n# (norm) specifies the type of parameter-free normalization.\n# (e.g. syncbatch, batch, instance)\n# (ks) specifies the size of kernel in the SPADE module (e.g. 3x3)\n# Example |config_text| will be spadesyncbatch3x3, or spadeinstance5x5.\n# Also, the other arguments are\n# |norm_nc|: the #channels of the normalized activations, hence the output dim of SPADE\n# |label_nc|: the #channels of the input semantic map, hence the input dim of SPADE\nclass SPADE(nn.Module):\n def __init__(self, config_text, norm_nc, label_nc):\n super().__init__()\n\n assert config_text.startswith('spade')\n parsed = re.search('spade(\\D+)(\\d)x\\d', config_text)\n param_free_norm_type = str(parsed.group(1))\n ks = int(parsed.group(2))\n\n if param_free_norm_type == 'instance':\n self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)\n elif param_free_norm_type == 'syncbatch':\n self.param_free_norm = SynchronizedBatchNorm2d(norm_nc, affine=False)\n elif param_free_norm_type == 'batch':\n self.param_free_norm = nn.BatchNorm2d(norm_nc, affine=False)\n else:\n raise ValueError('%s is not a recognized param-free norm type in SPADE'\n % param_free_norm_type)\n\n # The dimension of the intermediate embedding space. Yes, hardcoded.\n nhidden = 128\n\n pw = ks // 2\n self.mlp_shared = nn.Sequential(\n nn.Conv2d(label_nc, nhidden, kernel_size=ks, padding=pw),\n nn.ReLU()\n )\n self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw)\n self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw)\n\n def forward(self, x, segmap):\n\n # Part 1. generate parameter-free normalized activations\n normalized = self.param_free_norm(x)\n\n # Part 2. produce scaling and bias conditioned on semantic map\n segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest')\n actv = self.mlp_shared(segmap)\n gamma = self.mlp_gamma(actv)\n beta = self.mlp_beta(actv)\n\n # apply scale and bias\n out = normalized * (1 + gamma) + beta\n\n return out\n"
},
{
"path": "models/pix2pix_model.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nimport torch\nimport models.networks as networks\nimport util.util as util\n\n\nclass Pix2PixModel(torch.nn.Module):\n @staticmethod\n def modify_commandline_options(parser, is_train):\n networks.modify_commandline_options(parser, is_train)\n return parser\n\n def __init__(self, opt):\n super().__init__()\n self.opt = opt\n self.FloatTensor = torch.cuda.FloatTensor if self.use_gpu() \\\n else torch.FloatTensor\n self.ByteTensor = torch.cuda.ByteTensor if self.use_gpu() \\\n else torch.ByteTensor\n\n self.netG, self.netD, self.netE = self.initialize_networks(opt)\n\n # set loss functions\n if opt.isTrain:\n self.criterionGAN = networks.GANLoss(\n opt.gan_mode, tensor=self.FloatTensor, opt=self.opt)\n self.criterionFeat = torch.nn.L1Loss()\n if not opt.no_vgg_loss:\n self.criterionVGG = networks.VGGLoss(self.opt.gpu_ids)\n if opt.use_vae:\n self.KLDLoss = networks.KLDLoss()\n\n # Entry point for all calls involving forward pass\n # of deep networks. We used this approach since DataParallel module\n # can't parallelize custom functions, we branch to different\n # routines based on |mode|.\n def forward(self, data, mode):\n input_semantics, real_image = self.preprocess_input(data)\n\n if mode == 'generator':\n g_loss, generated = self.compute_generator_loss(\n input_semantics, real_image)\n return g_loss, generated\n elif mode == 'discriminator':\n d_loss = self.compute_discriminator_loss(\n input_semantics, real_image)\n return d_loss\n elif mode == 'encode_only':\n z, mu, logvar = self.encode_z(real_image)\n return mu, logvar\n elif mode == 'inference':\n with torch.no_grad():\n fake_image, _ = self.generate_fake(input_semantics, real_image)\n return fake_image\n else:\n raise ValueError(\"|mode| is invalid\")\n\n def create_optimizers(self, opt):\n G_params = list(self.netG.parameters())\n if opt.use_vae:\n G_params += list(self.netE.parameters())\n if opt.isTrain:\n D_params = list(self.netD.parameters())\n\n beta1, beta2 = opt.beta1, opt.beta2\n if opt.no_TTUR:\n G_lr, D_lr = opt.lr, opt.lr\n else:\n G_lr, D_lr = opt.lr / 2, opt.lr * 2\n\n optimizer_G = torch.optim.Adam(G_params, lr=G_lr, betas=(beta1, beta2))\n optimizer_D = torch.optim.Adam(D_params, lr=D_lr, betas=(beta1, beta2))\n\n return optimizer_G, optimizer_D\n\n def save(self, epoch):\n util.save_network(self.netG, 'G', epoch, self.opt)\n util.save_network(self.netD, 'D', epoch, self.opt)\n if self.opt.use_vae:\n util.save_network(self.netE, 'E', epoch, self.opt)\n\n ############################################################################\n # Private helper methods\n ############################################################################\n\n def initialize_networks(self, opt):\n netG = networks.define_G(opt)\n netD = networks.define_D(opt) if opt.isTrain else None\n netE = networks.define_E(opt) if opt.use_vae else None\n\n if not opt.isTrain or opt.continue_train:\n netG = util.load_network(netG, 'G', opt.which_epoch, opt)\n if opt.isTrain:\n netD = util.load_network(netD, 'D', opt.which_epoch, opt)\n if opt.use_vae:\n netE = util.load_network(netE, 'E', opt.which_epoch, opt)\n\n return netG, netD, netE\n\n # preprocess the input, such as moving the tensors to GPUs and\n # transforming the label map to one-hot encoding\n # |data|: dictionary of the input data\n\n def preprocess_input(self, data):\n # move to GPU and change data types\n data['label'] = data['label'].long()\n if self.use_gpu():\n data['label'] = data['label'].cuda()\n data['instance'] = data['instance'].cuda()\n data['image'] = data['image'].cuda()\n\n # create one-hot label map\n label_map = data['label']\n bs, _, h, w = label_map.size()\n nc = self.opt.label_nc + 1 if self.opt.contain_dontcare_label \\\n else self.opt.label_nc\n input_label = self.FloatTensor(bs, nc, h, w).zero_()\n input_semantics = input_label.scatter_(1, label_map, 1.0)\n\n # concatenate instance map if it exists\n if not self.opt.no_instance:\n inst_map = data['instance']\n instance_edge_map = self.get_edges(inst_map)\n input_semantics = torch.cat((input_semantics, instance_edge_map), dim=1)\n\n return input_semantics, data['image']\n\n def compute_generator_loss(self, input_semantics, real_image):\n G_losses = {}\n\n fake_image, KLD_loss = self.generate_fake(\n input_semantics, real_image, compute_kld_loss=self.opt.use_vae)\n\n if self.opt.use_vae:\n G_losses['KLD'] = KLD_loss\n\n pred_fake, pred_real = self.discriminate(\n input_semantics, fake_image, real_image)\n\n G_losses['GAN'] = self.criterionGAN(pred_fake, True,\n for_discriminator=False)\n\n if not self.opt.no_ganFeat_loss:\n num_D = len(pred_fake)\n GAN_Feat_loss = self.FloatTensor(1).fill_(0)\n for i in range(num_D): # for each discriminator\n # last output is the final prediction, so we exclude it\n num_intermediate_outputs = len(pred_fake[i]) - 1\n for j in range(num_intermediate_outputs): # for each layer output\n unweighted_loss = self.criterionFeat(\n pred_fake[i][j], pred_real[i][j].detach())\n GAN_Feat_loss += unweighted_loss * self.opt.lambda_feat / num_D\n G_losses['GAN_Feat'] = GAN_Feat_loss\n\n if not self.opt.no_vgg_loss:\n G_losses['VGG'] = self.criterionVGG(fake_image, real_image) \\\n * self.opt.lambda_vgg\n\n return G_losses, fake_image\n\n def compute_discriminator_loss(self, input_semantics, real_image):\n D_losses = {}\n with torch.no_grad():\n fake_image, _ = self.generate_fake(input_semantics, real_image)\n fake_image = fake_image.detach()\n fake_image.requires_grad_()\n\n pred_fake, pred_real = self.discriminate(\n input_semantics, fake_image, real_image)\n\n D_losses['D_Fake'] = self.criterionGAN(pred_fake, False,\n for_discriminator=True)\n D_losses['D_real'] = self.criterionGAN(pred_real, True,\n for_discriminator=True)\n\n return D_losses\n\n def encode_z(self, real_image):\n mu, logvar = self.netE(real_image)\n z = self.reparameterize(mu, logvar)\n return z, mu, logvar\n\n def generate_fake(self, input_semantics, real_image, compute_kld_loss=False):\n z = None\n KLD_loss = None\n if self.opt.use_vae:\n z, mu, logvar = self.encode_z(real_image)\n if compute_kld_loss:\n KLD_loss = self.KLDLoss(mu, logvar) * self.opt.lambda_kld\n\n fake_image = self.netG(input_semantics, z=z)\n\n assert (not compute_kld_loss) or self.opt.use_vae, \\\n \"You cannot compute KLD loss if opt.use_vae == False\"\n\n return fake_image, KLD_loss\n\n # Given fake and real image, return the prediction of discriminator\n # for each fake and real image.\n\n def discriminate(self, input_semantics, fake_image, real_image):\n fake_concat = torch.cat([input_semantics, fake_image], dim=1)\n real_concat = torch.cat([input_semantics, real_image], dim=1)\n\n # In Batch Normalization, the fake and real images are\n # recommended to be in the same batch to avoid disparate\n # statistics in fake and real images.\n # So both fake and real images are fed to D all at once.\n fake_and_real = torch.cat([fake_concat, real_concat], dim=0)\n\n discriminator_out = self.netD(fake_and_real)\n\n pred_fake, pred_real = self.divide_pred(discriminator_out)\n\n return pred_fake, pred_real\n\n # Take the prediction of fake and real images from the combined batch\n def divide_pred(self, pred):\n # the prediction contains the intermediate outputs of multiscale GAN,\n # so it's usually a list\n if type(pred) == list:\n fake = []\n real = []\n for p in pred:\n fake.append([tensor[:tensor.size(0) // 2] for tensor in p])\n real.append([tensor[tensor.size(0) // 2:] for tensor in p])\n else:\n fake = pred[:pred.size(0) // 2]\n real = pred[pred.size(0) // 2:]\n\n return fake, real\n\n def get_edges(self, t):\n edge = self.ByteTensor(t.size()).zero_()\n edge[:, :, :, 1:] = edge[:, :, :, 1:] | (t[:, :, :, 1:] != t[:, :, :, :-1])\n edge[:, :, :, :-1] = edge[:, :, :, :-1] | (t[:, :, :, 1:] != t[:, :, :, :-1])\n edge[:, :, 1:, :] = edge[:, :, 1:, :] | (t[:, :, 1:, :] != t[:, :, :-1, :])\n edge[:, :, :-1, :] = edge[:, :, :-1, :] | (t[:, :, 1:, :] != t[:, :, :-1, :])\n return edge.float()\n\n def reparameterize(self, mu, logvar):\n std = torch.exp(0.5 * logvar)\n eps = torch.randn_like(std)\n return eps.mul(std) + mu\n\n def use_gpu(self):\n return len(self.opt.gpu_ids) > 0\n"
},
{
"path": "options/__init__.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\""
},
{
"path": "options/base_options.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nimport sys\nimport argparse\nimport os\nfrom util import util\nimport torch\nimport models\nimport data\nimport pickle\n\n\nclass BaseOptions():\n def __init__(self):\n self.initialized = False\n\n def initialize(self, parser):\n # experiment specifics\n parser.add_argument('--name', type=str, default='label2coco', help='name of the experiment. It decides where to store samples and models')\n\n parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU')\n parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')\n parser.add_argument('--model', type=str, default='pix2pix', help='which model to use')\n parser.add_argument('--norm_G', type=str, default='spectralinstance', help='instance normalization or batch normalization')\n parser.add_argument('--norm_D', type=str, default='spectralinstance', help='instance normalization or batch normalization')\n parser.add_argument('--norm_E', type=str, default='spectralinstance', help='instance normalization or batch normalization')\n parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')\n\n # input/output sizes\n parser.add_argument('--batchSize', type=int, default=1, help='input batch size')\n parser.add_argument('--preprocess_mode', type=str, default='scale_width_and_crop', help='scaling and cropping of images at load time.', choices=(\"resize_and_crop\", \"crop\", \"scale_width\", \"scale_width_and_crop\", \"scale_shortside\", \"scale_shortside_and_crop\", \"fixed\", \"none\"))\n parser.add_argument('--load_size', type=int, default=1024, help='Scale images to this size. The final image will be cropped to --crop_size.')\n parser.add_argument('--crop_size', type=int, default=512, help='Crop to the width of crop_size (after initially scaling the images to load_size.)')\n parser.add_argument('--aspect_ratio', type=float, default=1.0, help='The ratio width/height. The final height of the load image will be crop_size/aspect_ratio')\n parser.add_argument('--label_nc', type=int, default=182, help='# of input label classes without unknown class. If you have unknown class as class label, specify --contain_dopntcare_label.')\n parser.add_argument('--contain_dontcare_label', action='store_true', help='if the label map contains dontcare label (dontcare=255)')\n parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels')\n\n # for setting inputs\n parser.add_argument('--dataroot', type=str, default='./datasets/cityscapes/')\n parser.add_argument('--dataset_mode', type=str, default='coco')\n parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')\n parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data argumentation')\n parser.add_argument('--nThreads', default=0, type=int, help='# threads for loading data')\n parser.add_argument('--max_dataset_size', type=int, default=sys.maxsize, help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')\n parser.add_argument('--load_from_opt_file', action='store_true', help='load the options from checkpoints and use that as default')\n parser.add_argument('--cache_filelist_write', action='store_true', help='saves the current filelist into a text file, so that it loads faster')\n parser.add_argument('--cache_filelist_read', action='store_true', help='reads from the file list cache')\n\n # for displays\n parser.add_argument('--display_winsize', type=int, default=400, help='display window size')\n\n # for generator\n parser.add_argument('--netG', type=str, default='spade', help='selects model to use for netG (pix2pixhd | spade)')\n parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer')\n parser.add_argument('--init_type', type=str, default='xavier', help='network initialization [normal|xavier|kaiming|orthogonal]')\n parser.add_argument('--init_variance', type=float, default=0.02, help='variance of the initialization distribution')\n parser.add_argument('--z_dim', type=int, default=256,\n help=\"dimension of the latent z vector\")\n\n # for instance-wise features\n parser.add_argument('--no_instance', action='store_true', help='if specified, do *not* add instance map as input')\n parser.add_argument('--nef', type=int, default=16, help='# of encoder filters in the first conv layer')\n parser.add_argument('--use_vae', action='store_true', help='enable training with an image encoder.')\n\n self.initialized = True\n return parser\n\n def gather_options(self):\n # initialize parser with basic options\n if not self.initialized:\n parser = argparse.ArgumentParser(\n formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n parser = self.initialize(parser)\n\n # get the basic options\n opt, unknown = parser.parse_known_args()\n\n # modify model-related parser options\n model_name = opt.model\n model_option_setter = models.get_option_setter(model_name)\n parser = model_option_setter(parser, self.isTrain)\n\n # modify dataset-related parser options\n dataset_mode = opt.dataset_mode\n dataset_option_setter = data.get_option_setter(dataset_mode)\n parser = dataset_option_setter(parser, self.isTrain)\n\n opt, unknown = parser.parse_known_args()\n\n # if there is opt_file, load it.\n # The previous default options will be overwritten\n if opt.load_from_opt_file:\n parser = self.update_options_from_file(parser, opt)\n\n opt = parser.parse_args()\n self.parser = parser\n return opt\n\n def print_options(self, opt):\n message = ''\n message += '----------------- Options ---------------\\n'\n for k, v in sorted(vars(opt).items()):\n comment = ''\n default = self.parser.get_default(k)\n if v != default:\n comment = '\\t[default: %s]' % str(default)\n message += '{:>25}: {:<30}{}\\n'.format(str(k), str(v), comment)\n message += '----------------- End -------------------'\n print(message)\n\n def option_file_path(self, opt, makedir=False):\n expr_dir = os.path.join(opt.checkpoints_dir, opt.name)\n if makedir:\n util.mkdirs(expr_dir)\n file_name = os.path.join(expr_dir, 'opt')\n return file_name\n\n def save_options(self, opt):\n file_name = self.option_file_path(opt, makedir=True)\n with open(file_name + '.txt', 'wt') as opt_file:\n for k, v in sorted(vars(opt).items()):\n comment = ''\n default = self.parser.get_default(k)\n if v != default:\n comment = '\\t[default: %s]' % str(default)\n opt_file.write('{:>25}: {:<30}{}\\n'.format(str(k), str(v), comment))\n\n with open(file_name + '.pkl', 'wb') as opt_file:\n pickle.dump(opt, opt_file)\n\n def update_options_from_file(self, parser, opt):\n new_opt = self.load_options(opt)\n for k, v in sorted(vars(opt).items()):\n if hasattr(new_opt, k) and v != getattr(new_opt, k):\n new_val = getattr(new_opt, k)\n parser.set_defaults(**{k: new_val})\n return parser\n\n def load_options(self, opt):\n file_name = self.option_file_path(opt, makedir=False)\n new_opt = pickle.load(open(file_name + '.pkl', 'rb'))\n return new_opt\n\n def parse(self, save=False):\n\n opt = self.gather_options()\n opt.isTrain = self.isTrain # train or test\n\n self.print_options(opt)\n if opt.isTrain:\n self.save_options(opt)\n\n # Set semantic_nc based on the option.\n # This will be convenient in many places\n opt.semantic_nc = opt.label_nc + \\\n (1 if opt.contain_dontcare_label else 0) + \\\n (0 if opt.no_instance else 1)\n\n # set gpu ids\n str_ids = opt.gpu_ids.split(',')\n opt.gpu_ids = []\n for str_id in str_ids:\n id = int(str_id)\n if id >= 0:\n opt.gpu_ids.append(id)\n if len(opt.gpu_ids) > 0:\n torch.cuda.set_device(opt.gpu_ids[0])\n\n assert len(opt.gpu_ids) == 0 or opt.batchSize % len(opt.gpu_ids) == 0, \\\n \"Batch size %d is wrong. It must be a multiple of # GPUs %d.\" \\\n % (opt.batchSize, len(opt.gpu_ids))\n\n self.opt = opt\n return self.opt\n"
},
{
"path": "options/test_options.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nfrom .base_options import BaseOptions\n\n\nclass TestOptions(BaseOptions):\n def initialize(self, parser):\n BaseOptions.initialize(self, parser)\n parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.')\n parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')\n parser.add_argument('--how_many', type=int, default=float(\"inf\"), help='how many test images to run')\n\n parser.set_defaults(preprocess_mode='scale_width_and_crop', crop_size=256, load_size=256, display_winsize=256)\n parser.set_defaults(serial_batches=True)\n parser.set_defaults(no_flip=True)\n parser.set_defaults(phase='test')\n self.isTrain = False\n return parser\n"
},
{
"path": "options/train_options.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nfrom .base_options import BaseOptions\n\n\nclass TrainOptions(BaseOptions):\n def initialize(self, parser):\n BaseOptions.initialize(self, parser)\n # for displays\n parser.add_argument('--display_freq', type=int, default=100, help='frequency of showing training results on screen')\n parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console')\n parser.add_argument('--save_latest_freq', type=int, default=5000, help='frequency of saving the latest results')\n parser.add_argument('--save_epoch_freq', type=int, default=10, help='frequency of saving checkpoints at the end of epochs')\n parser.add_argument('--no_html', action='store_true', help='do not save intermediate training results to [opt.checkpoints_dir]/[opt.name]/web/')\n parser.add_argument('--debug', action='store_true', help='only do one epoch and displays at each iteration')\n parser.add_argument('--tf_log', action='store_true', help='if specified, use tensorboard logging. Requires tensorflow installed')\n\n # for training\n parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')\n parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')\n parser.add_argument('--niter', type=int, default=50, help='# of iter at starting learning rate. This is NOT the total #epochs. Totla #epochs is niter + niter_decay')\n parser.add_argument('--niter_decay', type=int, default=0, help='# of iter to linearly decay learning rate to zero')\n parser.add_argument('--optimizer', type=str, default='adam')\n parser.add_argument('--beta1', type=float, default=0.0, help='momentum term of adam')\n parser.add_argument('--beta2', type=float, default=0.9, help='momentum term of adam')\n parser.add_argument('--no_TTUR', action='store_true', help='Use TTUR training scheme')\n\n # the default values for beta1 and beta2 differ by TTUR option\n opt, _ = parser.parse_known_args()\n if opt.no_TTUR:\n parser.set_defaults(beta1=0.5, beta2=0.999)\n\n parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')\n parser.add_argument('--D_steps_per_G', type=int, default=1, help='number of discriminator iterations per generator iterations.')\n\n # for discriminators\n parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in first conv layer')\n parser.add_argument('--lambda_feat', type=float, default=10.0, help='weight for feature matching loss')\n parser.add_argument('--lambda_vgg', type=float, default=10.0, help='weight for vgg loss')\n parser.add_argument('--no_ganFeat_loss', action='store_true', help='if specified, do *not* use discriminator feature matching loss')\n parser.add_argument('--no_vgg_loss', action='store_true', help='if specified, do *not* use VGG feature matching loss')\n parser.add_argument('--gan_mode', type=str, default='hinge', help='(ls|original|hinge)')\n parser.add_argument('--netD', type=str, default='multiscale', help='(n_layers|multiscale|image)')\n parser.add_argument('--lambda_kld', type=float, default=0.05)\n self.isTrain = True\n return parser\n"
},
{
"path": "requirements.txt",
"content": "torch>=1.0.0\ntorchvision\ndominate>=2.3.1\ndill\nscikit-image\n"
},
{
"path": "test.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nimport os\nfrom collections import OrderedDict\n\nimport data\nfrom options.test_options import TestOptions\nfrom models.pix2pix_model import Pix2PixModel\nfrom util.visualizer import Visualizer\nfrom util import html\n\nopt = TestOptions().parse()\n\ndataloader = data.create_dataloader(opt)\n\nmodel = Pix2PixModel(opt)\nmodel.eval()\n\nvisualizer = Visualizer(opt)\n\n# create a webpage that summarizes the all results\nweb_dir = os.path.join(opt.results_dir, opt.name,\n '%s_%s' % (opt.phase, opt.which_epoch))\nwebpage = html.HTML(web_dir,\n 'Experiment = %s, Phase = %s, Epoch = %s' %\n (opt.name, opt.phase, opt.which_epoch))\n\n# test\nfor i, data_i in enumerate(dataloader):\n if i * opt.batchSize >= opt.how_many:\n break\n\n generated = model(data_i, mode='inference')\n\n img_path = data_i['path']\n for b in range(generated.shape[0]):\n print('process image... %s' % img_path[b])\n visuals = OrderedDict([('input_label', data_i['label'][b]),\n ('synthesized_image', generated[b])])\n visualizer.save_images(webpage, visuals, img_path[b:b + 1])\n\nwebpage.save()\n"
},
{
"path": "train.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nimport sys\nfrom collections import OrderedDict\nfrom options.train_options import TrainOptions\nimport data\nfrom util.iter_counter import IterationCounter\nfrom util.visualizer import Visualizer\nfrom trainers.pix2pix_trainer import Pix2PixTrainer\n\n# parse options\nopt = TrainOptions().parse()\n\n# print options to help debugging\nprint(' '.join(sys.argv))\n\n# load the dataset\ndataloader = data.create_dataloader(opt)\n\n# create trainer for our model\ntrainer = Pix2PixTrainer(opt)\n\n# create tool for counting iterations\niter_counter = IterationCounter(opt, len(dataloader))\n\n# create tool for visualization\nvisualizer = Visualizer(opt)\n\nfor epoch in iter_counter.training_epochs():\n iter_counter.record_epoch_start(epoch)\n for i, data_i in enumerate(dataloader, start=iter_counter.epoch_iter):\n iter_counter.record_one_iteration()\n\n # Training\n # train generator\n if i % opt.D_steps_per_G == 0:\n trainer.run_generator_one_step(data_i)\n\n # train discriminator\n trainer.run_discriminator_one_step(data_i)\n\n # Visualizations\n if iter_counter.needs_printing():\n losses = trainer.get_latest_losses()\n visualizer.print_current_errors(epoch, iter_counter.epoch_iter,\n losses, iter_counter.time_per_iter)\n visualizer.plot_current_errors(losses, iter_counter.total_steps_so_far)\n\n if iter_counter.needs_displaying():\n visuals = OrderedDict([('input_label', data_i['label']),\n ('synthesized_image', trainer.get_latest_generated()),\n ('real_image', data_i['image'])])\n visualizer.display_current_results(visuals, epoch, iter_counter.total_steps_so_far)\n\n if iter_counter.needs_saving():\n print('saving the latest model (epoch %d, total_steps %d)' %\n (epoch, iter_counter.total_steps_so_far))\n trainer.save('latest')\n iter_counter.record_current_iter()\n\n trainer.update_learning_rate(epoch)\n iter_counter.record_epoch_end()\n\n if epoch % opt.save_epoch_freq == 0 or \\\n epoch == iter_counter.total_epochs:\n print('saving the model at the end of epoch %d, iters %d' %\n (epoch, iter_counter.total_steps_so_far))\n trainer.save('latest')\n trainer.save(epoch)\n\nprint('Training was successfully finished.')\n"
},
{
"path": "trainers/__init__.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n"
},
{
"path": "trainers/pix2pix_trainer.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nfrom models.networks.sync_batchnorm import DataParallelWithCallback\nfrom models.pix2pix_model import Pix2PixModel\n\n\nclass Pix2PixTrainer():\n \"\"\"\n Trainer creates the model and optimizers, and uses them to\n updates the weights of the network while reporting losses\n and the latest visuals to visualize the progress in training.\n \"\"\"\n\n def __init__(self, opt):\n self.opt = opt\n self.pix2pix_model = Pix2PixModel(opt)\n if len(opt.gpu_ids) > 0:\n self.pix2pix_model = DataParallelWithCallback(self.pix2pix_model,\n device_ids=opt.gpu_ids)\n self.pix2pix_model_on_one_gpu = self.pix2pix_model.module\n else:\n self.pix2pix_model_on_one_gpu = self.pix2pix_model\n\n self.generated = None\n if opt.isTrain:\n self.optimizer_G, self.optimizer_D = \\\n self.pix2pix_model_on_one_gpu.create_optimizers(opt)\n self.old_lr = opt.lr\n\n def run_generator_one_step(self, data):\n self.optimizer_G.zero_grad()\n g_losses, generated = self.pix2pix_model(data, mode='generator')\n g_loss = sum(g_losses.values()).mean()\n g_loss.backward()\n self.optimizer_G.step()\n self.g_losses = g_losses\n self.generated = generated\n\n def run_discriminator_one_step(self, data):\n self.optimizer_D.zero_grad()\n d_losses = self.pix2pix_model(data, mode='discriminator')\n d_loss = sum(d_losses.values()).mean()\n d_loss.backward()\n self.optimizer_D.step()\n self.d_losses = d_losses\n\n def get_latest_losses(self):\n return {**self.g_losses, **self.d_losses}\n\n def get_latest_generated(self):\n return self.generated\n\n def update_learning_rate(self, epoch):\n self.update_learning_rate(epoch)\n\n def save(self, epoch):\n self.pix2pix_model_on_one_gpu.save(epoch)\n\n ##################################################################\n # Helper functions\n ##################################################################\n\n def update_learning_rate(self, epoch):\n if epoch > self.opt.niter:\n lrd = self.opt.lr / self.opt.niter_decay\n new_lr = self.old_lr - lrd\n else:\n new_lr = self.old_lr\n\n if new_lr != self.old_lr:\n if self.opt.no_TTUR:\n new_lr_G = new_lr\n new_lr_D = new_lr\n else:\n new_lr_G = new_lr / 2\n new_lr_D = new_lr * 2\n\n for param_group in self.optimizer_D.param_groups:\n param_group['lr'] = new_lr_D\n for param_group in self.optimizer_G.param_groups:\n param_group['lr'] = new_lr_G\n print('update learning rate: %f -> %f' % (self.old_lr, new_lr))\n self.old_lr = new_lr\n"
},
{
"path": "util/__init__.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n"
},
{
"path": "util/coco.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\n\ndef id2label(id):\n if id == 182:\n id = 0\n else:\n id = id + 1\n labelmap = \\\n {0: 'unlabeled',\n 1: 'person',\n 2: 'bicycle',\n 3: 'car',\n 4: 'motorcycle',\n 5: 'airplane',\n 6: 'bus',\n 7: 'train',\n 8: 'truck',\n 9: 'boat',\n 10: 'traffic light',\n 11: 'fire hydrant',\n 12: 'street sign',\n 13: 'stop sign',\n 14: 'parking meter',\n 15: 'bench',\n 16: 'bird',\n 17: 'cat',\n 18: 'dog',\n 19: 'horse',\n 20: 'sheep',\n 21: 'cow',\n 22: 'elephant',\n 23: 'bear',\n 24: 'zebra',\n 25: 'giraffe',\n 26: 'hat',\n 27: 'backpack',\n 28: 'umbrella',\n 29: 'shoe',\n 30: 'eye glasses',\n 31: 'handbag',\n 32: 'tie',\n 33: 'suitcase',\n 34: 'frisbee',\n 35: 'skis',\n 36: 'snowboard',\n 37: 'sports ball',\n 38: 'kite',\n 39: 'baseball bat',\n 40: 'baseball glove',\n 41: 'skateboard',\n 42: 'surfboard',\n 43: 'tennis racket',\n 44: 'bottle',\n 45: 'plate',\n 46: 'wine glass',\n 47: 'cup',\n 48: 'fork',\n 49: 'knife',\n 50: 'spoon',\n 51: 'bowl',\n 52: 'banana',\n 53: 'apple',\n 54: 'sandwich',\n 55: 'orange',\n 56: 'broccoli',\n 57: 'carrot',\n 58: 'hot dog',\n 59: 'pizza',\n 60: 'donut',\n 61: 'cake',\n 62: 'chair',\n 63: 'couch',\n 64: 'potted plant',\n 65: 'bed',\n 66: 'mirror',\n 67: 'dining table',\n 68: 'window',\n 69: 'desk',\n 70: 'toilet',\n 71: 'door',\n 72: 'tv',\n 73: 'laptop',\n 74: 'mouse',\n 75: 'remote',\n 76: 'keyboard',\n 77: 'cell phone',\n 78: 'microwave',\n 79: 'oven',\n 80: 'toaster',\n 81: 'sink',\n 82: 'refrigerator',\n 83: 'blender',\n 84: 'book',\n 85: 'clock',\n 86: 'vase',\n 87: 'scissors',\n 88: 'teddy bear',\n 89: 'hair drier',\n 90: 'toothbrush',\n 91: 'hair brush', # Last class of Thing\n 92: 'banner', # Beginning of Stuff\n 93: 'blanket',\n 94: 'branch',\n 95: 'bridge',\n 96: 'building-other',\n 97: 'bush',\n 98: 'cabinet',\n 99: 'cage',\n 100: 'cardboard',\n 101: 'carpet',\n 102: 'ceiling-other',\n 103: 'ceiling-tile',\n 104: 'cloth',\n 105: 'clothes',\n 106: 'clouds',\n 107: 'counter',\n 108: 'cupboard',\n 109: 'curtain',\n 110: 'desk-stuff',\n 111: 'dirt',\n 112: 'door-stuff',\n 113: 'fence',\n 114: 'floor-marble',\n 115: 'floor-other',\n 116: 'floor-stone',\n 117: 'floor-tile',\n 118: 'floor-wood',\n 119: 'flower',\n 120: 'fog',\n 121: 'food-other',\n 122: 'fruit',\n 123: 'furniture-other',\n 124: 'grass',\n 125: 'gravel',\n 126: 'ground-other',\n 127: 'hill',\n 128: 'house',\n 129: 'leaves',\n 130: 'light',\n 131: 'mat',\n 132: 'metal',\n 133: 'mirror-stuff',\n 134: 'moss',\n 135: 'mountain',\n 136: 'mud',\n 137: 'napkin',\n 138: 'net',\n 139: 'paper',\n 140: 'pavement',\n 141: 'pillow',\n 142: 'plant-other',\n 143: 'plastic',\n 144: 'platform',\n 145: 'playingfield',\n 146: 'railing',\n 147: 'railroad',\n 148: 'river',\n 149: 'road',\n 150: 'rock',\n 151: 'roof',\n 152: 'rug',\n 153: 'salad',\n 154: 'sand',\n 155: 'sea',\n 156: 'shelf',\n 157: 'sky-other',\n 158: 'skyscraper',\n 159: 'snow',\n 160: 'solid-other',\n 161: 'stairs',\n 162: 'stone',\n 163: 'straw',\n 164: 'structural-other',\n 165: 'table',\n 166: 'tent',\n 167: 'textile-other',\n 168: 'towel',\n 169: 'tree',\n 170: 'vegetable',\n 171: 'wall-brick',\n 172: 'wall-concrete',\n 173: 'wall-other',\n 174: 'wall-panel',\n 175: 'wall-stone',\n 176: 'wall-tile',\n 177: 'wall-wood',\n 178: 'water-other',\n 179: 'waterdrops',\n 180: 'window-blind',\n 181: 'window-other',\n 182: 'wood'}\n if id in labelmap:\n return labelmap[id]\n else:\n return 'unknown'\n"
},
{
"path": "util/html.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nimport datetime\nimport dominate\nfrom dominate.tags import *\nimport os\n\n\nclass HTML:\n def __init__(self, web_dir, title, refresh=0):\n if web_dir.endswith('.html'):\n web_dir, html_name = os.path.split(web_dir)\n else:\n web_dir, html_name = web_dir, 'index.html'\n self.title = title\n self.web_dir = web_dir\n self.html_name = html_name\n self.img_dir = os.path.join(self.web_dir, 'images')\n if len(self.web_dir) > 0 and not os.path.exists(self.web_dir):\n os.makedirs(self.web_dir)\n if len(self.web_dir) > 0 and not os.path.exists(self.img_dir):\n os.makedirs(self.img_dir)\n\n self.doc = dominate.document(title=title)\n with self.doc:\n h1(datetime.datetime.now().strftime(\"%I:%M%p on %B %d, %Y\"))\n if refresh > 0:\n with self.doc.head:\n meta(http_equiv=\"refresh\", content=str(refresh))\n\n def get_image_dir(self):\n return self.img_dir\n\n def add_header(self, str):\n with self.doc:\n h3(str)\n\n def add_table(self, border=1):\n self.t = table(border=border, style=\"table-layout: fixed;\")\n self.doc.add(self.t)\n\n def add_images(self, ims, txts, links, width=512):\n self.add_table()\n with self.t:\n with tr():\n for im, txt, link in zip(ims, txts, links):\n with td(style=\"word-wrap: break-word;\", halign=\"center\", valign=\"top\"):\n with p():\n with a(href=os.path.join('images', link)):\n img(style=\"width:%dpx\" % (width), src=os.path.join('images', im))\n br()\n p(txt.encode('utf-8'))\n\n def save(self):\n html_file = os.path.join(self.web_dir, self.html_name)\n f = open(html_file, 'wt')\n f.write(self.doc.render())\n f.close()\n\n\nif __name__ == '__main__':\n html = HTML('web/', 'test_html')\n html.add_header('hello world')\n\n ims = []\n txts = []\n links = []\n for n in range(4):\n ims.append('image_%d.jpg' % n)\n txts.append('text_%d' % n)\n links.append('image_%d.jpg' % n)\n html.add_images(ims, txts, links)\n html.save()\n"
},
{
"path": "util/iter_counter.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nimport os\nimport time\nimport numpy as np\n\n\n# Helper class that keeps track of training iterations\nclass IterationCounter():\n def __init__(self, opt, dataset_size):\n self.opt = opt\n self.dataset_size = dataset_size\n\n self.first_epoch = 1\n self.total_epochs = opt.niter + opt.niter_decay\n self.epoch_iter = 0 # iter number within each epoch\n self.iter_record_path = os.path.join(self.opt.checkpoints_dir, self.opt.name, 'iter.txt')\n if opt.isTrain and opt.continue_train:\n try:\n self.first_epoch, self.epoch_iter = np.loadtxt(\n self.iter_record_path, delimiter=',', dtype=int)\n print('Resuming from epoch %d at iteration %d' % (self.first_epoch, self.epoch_iter))\n except:\n print('Could not load iteration record at %s. Starting from beginning.' %\n self.iter_record_path)\n\n self.total_steps_so_far = (self.first_epoch - 1) * dataset_size + self.epoch_iter\n\n # return the iterator of epochs for the training\n def training_epochs(self):\n return range(self.first_epoch, self.total_epochs + 1)\n\n def record_epoch_start(self, epoch):\n self.epoch_start_time = time.time()\n self.epoch_iter = 0\n self.last_iter_time = time.time()\n self.current_epoch = epoch\n\n def record_one_iteration(self):\n current_time = time.time()\n\n # the last remaining batch is dropped (see data/__init__.py),\n # so we can assume batch size is always opt.batchSize\n self.time_per_iter = (current_time - self.last_iter_time) / self.opt.batchSize\n self.last_iter_time = current_time\n self.total_steps_so_far += self.opt.batchSize\n self.epoch_iter += self.opt.batchSize\n\n def record_epoch_end(self):\n current_time = time.time()\n self.time_per_epoch = current_time - self.epoch_start_time\n print('End of epoch %d / %d \\t Time Taken: %d sec' %\n (self.current_epoch, self.total_epochs, self.time_per_epoch))\n if self.current_epoch % self.opt.save_epoch_freq == 0:\n np.savetxt(self.iter_record_path, (self.current_epoch + 1, 0),\n delimiter=',', fmt='%d')\n print('Saved current iteration count at %s.' % self.iter_record_path)\n\n def record_current_iter(self):\n np.savetxt(self.iter_record_path, (self.current_epoch, self.epoch_iter),\n delimiter=',', fmt='%d')\n print('Saved current iteration count at %s.' % self.iter_record_path)\n\n def needs_saving(self):\n return (self.total_steps_so_far % self.opt.save_latest_freq) < self.opt.batchSize\n\n def needs_printing(self):\n return (self.total_steps_so_far % self.opt.print_freq) < self.opt.batchSize\n\n def needs_displaying(self):\n return (self.total_steps_so_far % self.opt.display_freq) < self.opt.batchSize\n"
},
{
"path": "util/util.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nimport re\nimport importlib\nimport torch\nfrom argparse import Namespace\nimport numpy as np\nfrom PIL import Image\nimport os\nimport argparse\nimport dill as pickle\nimport util.coco\n\n\ndef save_obj(obj, name):\n with open(name, 'wb') as f:\n pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)\n\n\ndef load_obj(name):\n with open(name, 'rb') as f:\n return pickle.load(f)\n\n# returns a configuration for creating a generator\n# |default_opt| should be the opt of the current experiment\n# |**kwargs|: if any configuration should be overriden, it can be specified here\n\n\ndef copyconf(default_opt, **kwargs):\n conf = argparse.Namespace(**vars(default_opt))\n for key in kwargs:\n print(key, kwargs[key])\n setattr(conf, key, kwargs[key])\n return conf\n\n\ndef tile_images(imgs, picturesPerRow=4):\n \"\"\" Code borrowed from\n https://stackoverflow.com/questions/26521365/cleanly-tile-numpy-array-of-images-stored-in-a-flattened-1d-format/26521997\n \"\"\"\n\n # Padding\n if imgs.shape[0] % picturesPerRow == 0:\n rowPadding = 0\n else:\n rowPadding = picturesPerRow - imgs.shape[0] % picturesPerRow\n if rowPadding > 0:\n imgs = np.concatenate([imgs, np.zeros((rowPadding, *imgs.shape[1:]), dtype=imgs.dtype)], axis=0)\n\n # Tiling Loop (The conditionals are not necessary anymore)\n tiled = []\n for i in range(0, imgs.shape[0], picturesPerRow):\n tiled.append(np.concatenate([imgs[j] for j in range(i, i + picturesPerRow)], axis=1))\n\n tiled = np.concatenate(tiled, axis=0)\n return tiled\n\n\n# Converts a Tensor into a Numpy array\n# |imtype|: the desired type of the converted numpy array\ndef tensor2im(image_tensor, imtype=np.uint8, normalize=True, tile=False):\n if isinstance(image_tensor, list):\n image_numpy = []\n for i in range(len(image_tensor)):\n image_numpy.append(tensor2im(image_tensor[i], imtype, normalize))\n return image_numpy\n\n if image_tensor.dim() == 4:\n # transform each image in the batch\n images_np = []\n for b in range(image_tensor.size(0)):\n one_image = image_tensor[b]\n one_image_np = tensor2im(one_image)\n images_np.append(one_image_np.reshape(1, *one_image_np.shape))\n images_np = np.concatenate(images_np, axis=0)\n if tile:\n images_tiled = tile_images(images_np)\n return images_tiled\n else:\n return images_np\n\n if image_tensor.dim() == 2:\n image_tensor = image_tensor.unsqueeze(0)\n image_numpy = image_tensor.detach().cpu().float().numpy()\n if normalize:\n image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0\n else:\n image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0\n image_numpy = np.clip(image_numpy, 0, 255)\n if image_numpy.shape[2] == 1:\n image_numpy = image_numpy[:, :, 0]\n return image_numpy.astype(imtype)\n\n\n# Converts a one-hot tensor into a colorful label map\ndef tensor2label(label_tensor, n_label, imtype=np.uint8, tile=False):\n if label_tensor.dim() == 4:\n # transform each image in the batch\n images_np = []\n for b in range(label_tensor.size(0)):\n one_image = label_tensor[b]\n one_image_np = tensor2label(one_image, n_label, imtype)\n images_np.append(one_image_np.reshape(1, *one_image_np.shape))\n images_np = np.concatenate(images_np, axis=0)\n if tile:\n images_tiled = tile_images(images_np)\n return images_tiled\n else:\n images_np = images_np[0]\n return images_np\n\n if label_tensor.dim() == 1:\n return np.zeros((64, 64, 3), dtype=np.uint8)\n if n_label == 0:\n return tensor2im(label_tensor, imtype)\n label_tensor = label_tensor.cpu().float()\n if label_tensor.size()[0] > 1:\n label_tensor = label_tensor.max(0, keepdim=True)[1]\n label_tensor = Colorize(n_label)(label_tensor)\n label_numpy = np.transpose(label_tensor.numpy(), (1, 2, 0))\n result = label_numpy.astype(imtype)\n return result\n\n\ndef save_image(image_numpy, image_path, create_dir=False):\n if create_dir:\n os.makedirs(os.path.dirname(image_path), exist_ok=True)\n if len(image_numpy.shape) == 2:\n image_numpy = np.expand_dims(image_numpy, axis=2)\n if image_numpy.shape[2] == 1:\n image_numpy = np.repeat(image_numpy, 3, 2)\n image_pil = Image.fromarray(image_numpy)\n\n # save to png\n image_pil.save(image_path.replace('.jpg', '.png'))\n\n\ndef mkdirs(paths):\n if isinstance(paths, list) and not isinstance(paths, str):\n for path in paths:\n mkdir(path)\n else:\n mkdir(paths)\n\n\ndef mkdir(path):\n if not os.path.exists(path):\n os.makedirs(path)\n\n\ndef atoi(text):\n return int(text) if text.isdigit() else text\n\n\ndef natural_keys(text):\n '''\n alist.sort(key=natural_keys) sorts in human order\n http://nedbatchelder.com/blog/200712/human_sorting.html\n (See Toothy's implementation in the comments)\n '''\n return [atoi(c) for c in re.split('(\\d+)', text)]\n\n\ndef natural_sort(items):\n items.sort(key=natural_keys)\n\n\ndef str2bool(v):\n if v.lower() in ('yes', 'true', 't', 'y', '1'):\n return True\n elif v.lower() in ('no', 'false', 'f', 'n', '0'):\n return False\n else:\n raise argparse.ArgumentTypeError('Boolean value expected.')\n\n\ndef find_class_in_module(target_cls_name, module):\n target_cls_name = target_cls_name.replace('_', '').lower()\n clslib = importlib.import_module(module)\n cls = None\n for name, clsobj in clslib.__dict__.items():\n if name.lower() == target_cls_name:\n cls = clsobj\n\n if cls is None:\n print(\"In %s, there should be a class whose name matches %s in lowercase without underscore(_)\" % (module, target_cls_name))\n exit(0)\n\n return cls\n\n\ndef save_network(net, label, epoch, opt):\n save_filename = '%s_net_%s.pth' % (epoch, label)\n save_path = os.path.join(opt.checkpoints_dir, opt.name, save_filename)\n torch.save(net.cpu().state_dict(), save_path)\n if len(opt.gpu_ids) and torch.cuda.is_available():\n net.cuda()\n\n\ndef load_network(net, label, epoch, opt):\n save_filename = '%s_net_%s.pth' % (epoch, label)\n save_dir = os.path.join(opt.checkpoints_dir, opt.name)\n save_path = os.path.join(save_dir, save_filename)\n weights = torch.load(save_path)\n net.load_state_dict(weights)\n return net\n\n\n###############################################################################\n# Code from\n# https://github.com/ycszen/pytorch-seg/blob/master/transform.py\n# Modified so it complies with the Citscape label map colors\n###############################################################################\ndef uint82bin(n, count=8):\n \"\"\"returns the binary of integer n, count refers to amount of bits\"\"\"\n return ''.join([str((n >> y) & 1) for y in range(count - 1, -1, -1)])\n\n\ndef labelcolormap(N):\n if N == 35: # cityscape\n cmap = np.array([(0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (111, 74, 0), (81, 0, 81),\n (128, 64, 128), (244, 35, 232), (250, 170, 160), (230, 150, 140), (70, 70, 70), (102, 102, 156), (190, 153, 153),\n (180, 165, 180), (150, 100, 100), (150, 120, 90), (153, 153, 153), (153, 153, 153), (250, 170, 30), (220, 220, 0),\n (107, 142, 35), (152, 251, 152), (70, 130, 180), (220, 20, 60), (255, 0, 0), (0, 0, 142), (0, 0, 70),\n (0, 60, 100), (0, 0, 90), (0, 0, 110), (0, 80, 100), (0, 0, 230), (119, 11, 32), (0, 0, 142)],\n dtype=np.uint8)\n else:\n cmap = np.zeros((N, 3), dtype=np.uint8)\n for i in range(N):\n r, g, b = 0, 0, 0\n id = i + 1 # let's give 0 a color\n for j in range(7):\n str_id = uint82bin(id)\n r = r ^ (np.uint8(str_id[-1]) << (7 - j))\n g = g ^ (np.uint8(str_id[-2]) << (7 - j))\n b = b ^ (np.uint8(str_id[-3]) << (7 - j))\n id = id >> 3\n cmap[i, 0] = r\n cmap[i, 1] = g\n cmap[i, 2] = b\n\n if N == 182: # COCO\n important_colors = {\n 'sea': (54, 62, 167),\n 'sky-other': (95, 219, 255),\n 'tree': (140, 104, 47),\n 'clouds': (170, 170, 170),\n 'grass': (29, 195, 49)\n }\n for i in range(N):\n name = util.coco.id2label(i)\n if name in important_colors:\n color = important_colors[name]\n cmap[i] = np.array(list(color))\n\n return cmap\n\n\nclass Colorize(object):\n def __init__(self, n=35):\n self.cmap = labelcolormap(n)\n self.cmap = torch.from_numpy(self.cmap[:n])\n\n def __call__(self, gray_image):\n size = gray_image.size()\n color_image = torch.ByteTensor(3, size[1], size[2]).fill_(0)\n\n for label in range(0, len(self.cmap)):\n mask = (label == gray_image[0]).cpu()\n color_image[0][mask] = self.cmap[label][0]\n color_image[1][mask] = self.cmap[label][1]\n color_image[2][mask] = self.cmap[label][2]\n\n return color_image\n"
},
{
"path": "util/visualizer.py",
"content": "\"\"\"\nCopyright (C) 2019 NVIDIA Corporation. All rights reserved.\nLicensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).\n\"\"\"\n\nimport os\nimport ntpath\nimport time\nfrom . import util\nfrom . import html\nimport scipy.misc\ntry:\n from StringIO import StringIO # Python 2.7\nexcept ImportError:\n from io import BytesIO # Python 3.x\n\nclass Visualizer():\n def __init__(self, opt):\n self.opt = opt\n self.tf_log = opt.isTrain and opt.tf_log\n self.use_html = opt.isTrain and not opt.no_html\n self.win_size = opt.display_winsize\n self.name = opt.name\n if self.tf_log:\n import tensorflow as tf\n self.tf = tf\n self.log_dir = os.path.join(opt.checkpoints_dir, opt.name, 'logs')\n self.writer = tf.summary.FileWriter(self.log_dir)\n\n if self.use_html:\n self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web')\n self.img_dir = os.path.join(self.web_dir, 'images')\n print('create web directory %s...' % self.web_dir)\n util.mkdirs([self.web_dir, self.img_dir])\n if opt.isTrain:\n self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')\n with open(self.log_name, \"a\") as log_file:\n now = time.strftime(\"%c\")\n log_file.write('================ Training Loss (%s) ================\\n' % now)\n\n # |visuals|: dictionary of images to display or save\n def display_current_results(self, visuals, epoch, step):\n\n ## convert tensors to numpy arrays\n visuals = self.convert_visuals_to_numpy(visuals)\n \n if self.tf_log: # show images in tensorboard output\n img_summaries = []\n for label, image_numpy in visuals.items():\n # Write the image to a string\n try:\n s = StringIO()\n except:\n s = BytesIO()\n if len(image_numpy.shape) >= 4:\n image_numpy = image_numpy[0]\n scipy.misc.toimage(image_numpy).save(s, format=\"jpeg\")\n # Create an Image object\n img_sum = self.tf.Summary.Image(encoded_image_string=s.getvalue(), height=image_numpy.shape[0], width=image_numpy.shape[1])\n # Create a Summary value\n img_summaries.append(self.tf.Summary.Value(tag=label, image=img_sum))\n\n # Create and write Summary\n summary = self.tf.Summary(value=img_summaries)\n self.writer.add_summary(summary, step)\n\n if self.use_html: # save images to a html file\n for label, image_numpy in visuals.items():\n if isinstance(image_numpy, list):\n for i in range(len(image_numpy)):\n img_path = os.path.join(self.img_dir, 'epoch%.3d_iter%.3d_%s_%d.png' % (epoch, step, label, i))\n util.save_image(image_numpy[i], img_path)\n else:\n img_path = os.path.join(self.img_dir, 'epoch%.3d_iter%.3d_%s.png' % (epoch, step, label))\n if len(image_numpy.shape) >= 4:\n image_numpy = image_numpy[0] \n util.save_image(image_numpy, img_path)\n\n # update website\n webpage = html.HTML(self.web_dir, 'Experiment name = %s' % self.name, refresh=5)\n for n in range(epoch, 0, -1):\n webpage.add_header('epoch [%d]' % n)\n ims = []\n txts = []\n links = []\n\n for label, image_numpy in visuals.items():\n if isinstance(image_numpy, list):\n for i in range(len(image_numpy)):\n img_path = 'epoch%.3d_iter%.3d_%s_%d.png' % (n, step, label, i)\n ims.append(img_path)\n txts.append(label+str(i))\n links.append(img_path)\n else:\n img_path = 'epoch%.3d_iter%.3d_%s.png' % (n, step, label)\n ims.append(img_path)\n txts.append(label)\n links.append(img_path)\n if len(ims) < 10:\n webpage.add_images(ims, txts, links, width=self.win_size)\n else:\n num = int(round(len(ims)/2.0))\n webpage.add_images(ims[:num], txts[:num], links[:num], width=self.win_size)\n webpage.add_images(ims[num:], txts[num:], links[num:], width=self.win_size)\n webpage.save()\n\n # errors: dictionary of error labels and values\n def plot_current_errors(self, errors, step):\n if self.tf_log:\n for tag, value in errors.items():\n value = value.mean().float()\n summary = self.tf.Summary(value=[self.tf.Summary.Value(tag=tag, simple_value=value)])\n self.writer.add_summary(summary, step)\n\n # errors: same format as |errors| of plotCurrentErrors\n def print_current_errors(self, epoch, i, errors, t):\n message = '(epoch: %d, iters: %d, time: %.3f) ' % (epoch, i, t)\n for k, v in errors.items():\n #print(v)\n #if v != 0:\n v = v.mean().float()\n message += '%s: %.3f ' % (k, v)\n\n print(message)\n with open(self.log_name, \"a\") as log_file:\n log_file.write('%s\\n' % message)\n\n def convert_visuals_to_numpy(self, visuals):\n for key, t in visuals.items():\n tile = self.opt.batchSize > 8\n if 'input_label' == key:\n t = util.tensor2label(t, self.opt.label_nc + 2, tile=tile)\n else:\n t = util.tensor2im(t, tile=tile)\n visuals[key] = t\n return visuals\n\n # save image to the disk\n def save_images(self, webpage, visuals, image_path): \n visuals = self.convert_visuals_to_numpy(visuals) \n \n image_dir = webpage.get_image_dir()\n short_path = ntpath.basename(image_path[0])\n name = os.path.splitext(short_path)[0]\n\n webpage.add_header(name)\n ims = []\n txts = []\n links = []\n\n for label, image_numpy in visuals.items():\n image_name = os.path.join(label, '%s.png' % (name))\n save_path = os.path.join(image_dir, image_name)\n util.save_image(image_numpy, save_path, create_dir=True)\n\n ims.append(image_name)\n txts.append(label)\n links.append(image_name)\n webpage.add_images(ims, txts, links, width=self.win_size)\n"
}
]
Semantic Image Synthesis with Spatially-Adaptive Normalization
\n\tTaesung Park \n\tMing-Yu Liu \n\tTing-Chun Wang \n\tJun-Yan Zhu \n\t
\n\t\n\t\tUC Berkeley \n\t\tNVIDIA \n\t\tMIT \n\t
in CVPR 2019 (Oral)
Paper | Code | Online Demo App
\n\n
\n\t\t\t![](\"images/treepond.gif\") \n\t\t | \n\t\t\n\t\t\t![](\"images/ocean.gif\") \n\t\t | \n\t
Abstract
\n\nWe propose spatially-adaptive normalization, a simple but effective layer for synthesizing photorealistic images given an input semantic layout. Previous methods directly feed the semantic layout as input to the network, which is then processed through stacks of convolution, normalization, and nonlinearity layers. We show that this is suboptimal because the normalization layers tend to wash away semantic information. To address the issue, we propose using the input layout for modulating the activations in normalization layers through a spatially-adaptive, learned transformation. Experiments on several challenging datasets demonstrate the advantage of the proposed method compared to existing approaches, regarding both visual fidelity and alignment with input layouts. Finally, our model allows users to easily control the style and content of synthesis results as well as create multi-modal results.
\n\n\n\nPaper
\narxiv, 2019.
\n\n\n\nCitation
\nTaesung Park, Ming-Yu Liu, Ting-Chun Wang, and Jun-Yan Zhu.
\"Semantic Image Synthesis with Spatially-Adaptive Normalization\", in CVPR, 2019.\nBibtex\n\n\n
Code
\n\n\n\n\n
\n\t\tVideo of Interactive Demo App (GauGAN)\n\t\t\n\t | \n\t\n\t\n\t\tIntroduction of SPADE at GTC 2019\n\t\t\n\t | \n\n\t
\n
Brief Description of the Method
\n\n
In many common normalization techniques such as Batch Normalization (Ioffe et al., 2015), there are learned affine layers (as in PyTorch and TensorFlow) that are applied after the actual normalization step. In SPADE, the affine layer is learned from semantic segmentation map. This is similar to Conditional Normalization (De Vries et al., 2017 and Dumoulin et al., 2016), except that the learned affine parameters now need to be spatially-adaptive, which means we will use different scaling and bias for each semantic label. Using this simple method, semantic signal can act on all layer outputs, unaffected by the normalization process which may lose such information. Moreover, because the semantic information is provided via SPADE layers, random latent vector may be used as input to the network, which can be used to manipulate the style of the generated images. \n
\n\n\n\n
Comparison to Existing Methods
\nSPADE outperforms existing methods on the COCO-Stuff dataset, which is more challenging than the Cityscapes dataset due to more diverse scenes and labels. The images above are the ones authors liked. \n
\n\n\n
\n
Applying on Flickr Images
\nSince SPADE works on diverse labels, it can be trained with an existing semantic segmentation network to learn the reverse mapping from semantic maps to photos. These images were generated from SPADE trained on 40k images scraped from Flickr.\n
\n\n\n
Code and Trained Models
\n\tPlease visit our github repo.
\n\n\n
Online Demo
\n\tWe released an online demo of GauGAN, our interactive app that generates realistic landscape images from the layout users draw. The model was trained on landscape images scraped from Flickr.com. We released an online demo that has the same features. Please visit our online demo page.
\n\n\n\n
Acknowledgement
\nWe thank Alyosha Efros and Jan Kautz for insightful advice. Taesung Park contributed to the work during his internship at NVIDIA. His Ph.D. is supported by Samsung Scholarship.
\n\n\n