def define_D(input_nc,
ndf,
netD,
n_layers_D=3,
norm='batch',
init_type='normal',
init_gain=0.02):
"""Create a discriminator
Parameters:
input_nc (int) -- the number of channels in input images
ndf (int) -- the number of filters in the first conv layer
netD (str) -- the architecture's name: basic | n_layers | pixel
n_layers_D (int) -- the number of conv layers in the discriminator; effective when netD=='n_layers'
norm (str) -- the type of normalization layers used in the network.
init_type (str) -- the name of the initialization method.
init_gain (float) -- scaling factor for normal, xavier and orthogonal.
gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
Returns a discriminator
Our current implementation provides three types of discriminators:
[basic]: 'PatchGAN' classifier described in the original pix2pix paper.
It can classify whether 70×70 overlapping patches are real or fake.
Such a patch-level discriminator architecture has fewer parameters
than a full-image discriminator and can work on arbitrarily-sized images
in a fully convolutional fashion.
[n_layers]: With this mode, you cna specify the number of conv layers in the discriminator
with the parameter <n_layers_D> (default=3 as used in [basic] (PatchGAN).)
[pixel]: 1x1 PixelGAN discriminator can classify whether a pixel is real or not.
It encourages greater color diversity but has no effect on spatial statistics.
The discriminator has been initialized by <init_net>. It uses Leakly RELU for non-linearity.
"""
net = None
norm_layer = get_norm_layer(norm_type=norm)
if netD == 'basic': # default PatchGAN classifier
net = NLayerDiscriminator(input_nc,
ndf,
n_layers=3,
norm_layer=norm_layer)
elif netD == 'n_layers': # more options
net = NLayerDiscriminator(input_nc,
ndf,
n_layers_D,
norm_layer=norm_layer)
elif netD == 'pixel': # classify if each pixel is real or fake
net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer)
else:
raise NotImplementedError(
'Discriminator model name [%s] is not recognized' % netD)
init_weights(net, init_type, init_gain)
return net
def __init__(
self,
texture_channels=3,
cloth_channels=19,
num_roi=12,
norm_type="batch",
dropout=0.5,
unet_type="pix2pix",
img_size=128,
):
super(TextureModule, self).__init__()
self.roi_align = ROIAlign(output_size=(128, 128),
spatial_scale=1,
sampling_ratio=1)
self.num_roi = num_roi
channels = texture_channels * num_roi
self.encode = UNetDown(channels, channels)
# UNET
if unet_type == "pix2pix":
# fast log2 of img_size, int only. E.g. if size=128 => num_downs=7
num_downs = math.frexp(img_size)[1] - 1
use_dropout = True if dropout is not None else False
norm_layer = get_norm_layer(norm_type=norm_type)
self.unet = pix2pix_modules.UnetGenerator(
channels + cloth_channels,
texture_channels,
num_downs,
norm_layer=norm_layer,
use_dropout=use_dropout,
)
else:
self.unet = nn.Sequential(
UNetDown(channels + cloth_channels, 64, normalize=False),
UNetDown(64, 128),
UNetDown(128, 256),
UNetDown(256, 512, dropout=dropout),
UNetDown(512, 1024, dropout=dropout),
UNetDown(1024, 1024, normalize=False, dropout=dropout),
UNetUp(1024, 1024, dropout=dropout),
UNetUp(2 * 1024, 512, dropout=dropout),
UNetUp(2 * 512, 256),
UNetUp(2 * 256, 128),
UNetUp(2 * 128, 64),
# upsample and pad
nn.Upsample(scale_factor=2),
nn.ZeroPad2d((1, 0, 1, 0)),
nn.Conv2d(128, texture_channels, 4, padding=1),
nn.Tanh(),
)
def define_G(self):
if self.opt.netG == "unet_128":
norm_layer = get_norm_layer("batch")
return UnetGenerator(self.opt.texture_channels,
self.opt.texture_channels,
7,
64,
norm_layer=norm_layer,
use_dropout=True)
elif self.opt.netG == "swapnet":
return TextureModule(
texture_channels=self.opt.texture_channels,
cloth_channels=self.opt.cloth_channels,
num_roi=self.opt.body_channels,
img_size=self.opt.crop_size,
norm_type=self.opt.norm,
)
else:
raise ValueError("Cannot find implementation for " + self.opt.netG)
def define_G(
input_nc,
output_nc,
ngf,
netG,
norm="batch",
use_dropout=False,
init_type="normal",
init_gain=0.02,
):
"""Create a generator
Parameters:
input_nc (int) -- the number of channels in input images
output_nc (int) -- the number of channels in output images
ngf (int) -- the number of filters in the last conv layer
netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128
norm (str) -- the name of normalization layers used in the network: batch | instance | none
use_dropout (bool) -- if use dropout layers.
init_type (str) -- the name of our initialization method.
init_gain (float) -- scaling factor for normal, xavier and orthogonal.
gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
Returns a generator
Our current implementation provides two types of generators:
U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images)
The original U-Net paper: https://arxiv.org/abs/1505.04597
Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks)
Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations.
We adapt Torch code from Justin Johnson's neural style transfer project (https://github.com/jcjohnson/fast-neural-style).
The generator has been initialized by <init_net>. It uses RELU for non-linearity.
"""
net = None
norm_layer = get_norm_layer(norm_type=norm)
if netG == "resnet_9blocks":
raise NotImplementedError
elif netG == "resnet_6blocks":
raise NotImplementedError
elif netG == "unet_128":
net = UnetGenerator(input_nc,
output_nc,
7,
ngf,
norm_layer=norm_layer,
use_dropout=use_dropout)
elif netG == "unet_256":
net = UnetGenerator(input_nc,
output_nc,
8,
ngf,
norm_layer=norm_layer,
use_dropout=use_dropout)
else:
raise NotImplementedError(
"Generator model name [%s] is not recognized" % netG)
init_weights(net, init_type, init_gain)
return net
