diff --git a/.gitignore b/.gitignore index 822f94d6b34e4d0ae0f53420a3502105e9492763..5ebe6767b22921fcdb18a9fc3662b24db6ec025c 100644 --- a/.gitignore +++ b/.gitignore @@ -143,6 +143,7 @@ video_tmp/ result/ nohup.out #./ +/.vscode /pix2pix /pix2pixHD /tmp diff --git a/models/BVDNet.py b/models/BVDNet.py index ab5adeb38a87e27d74dbaed2acf6daf7124920be..7734e866e5e97f467f9f700775449384c93803c8 100644 --- a/models/BVDNet.py +++ b/models/BVDNet.py @@ -4,17 +4,38 @@ import torch.nn.functional as F from .pix2pixHD_model import * from .model_util import * +class UpBlock(nn.Module): + def __init__(self, in_channel, out_channel, kernel_size=3, padding=1): + super().__init__() + + self.convup = nn.Sequential( + nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False), + nn.ReflectionPad2d(padding), + # EqualConv2d(out_channel, out_channel, kernel_size, padding=padding), + SpectralNorm(nn.Conv2d(in_channel, out_channel, kernel_size)), + nn.LeakyReLU(0.2), + # Blur(out_channel), + ) + + + def forward(self, input): + + outup = self.convup(input) + + return outup + class Encoder2d(nn.Module): - def __init__(self, input_nc, ngf=64, n_downsampling=3, norm_layer=nn.BatchNorm2d, activation = nn.ReLU(True)): + def __init__(self, input_nc, ngf=64, n_downsampling=3, activation = nn.LeakyReLU(0.2)): super(Encoder2d, self).__init__() - model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation] + model = [nn.ReflectionPad2d(3), SpectralNorm(nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0)), activation] ### downsample for i in range(n_downsampling): mult = 2**i - model += [nn.ReflectionPad2d(1),nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=0), - norm_layer(ngf * mult * 2), activation] + model += [ nn.ReflectionPad2d(1), + SpectralNorm(nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=0)), + activation] self.model = nn.Sequential(*model) @@ -22,15 +43,15 @@ class Encoder2d(nn.Module): return self.model(input) class Encoder3d(nn.Module): - def __init__(self, input_nc, ngf=64, n_downsampling=3, norm_layer=nn.BatchNorm3d,activation = nn.ReLU(True)): + def __init__(self, input_nc, ngf=64, n_downsampling=3, activation = nn.LeakyReLU(0.2)): super(Encoder3d, self).__init__() - model = [nn.Conv3d(input_nc, ngf, kernel_size=3, padding=1), norm_layer(ngf), activation] + model = [SpectralNorm(nn.Conv3d(input_nc, ngf, kernel_size=3, padding=1)), activation] ### downsample for i in range(n_downsampling): mult = 2**i - model += [nn.Conv3d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1), - norm_layer(ngf * mult * 2), activation] + model += [ SpectralNorm(nn.Conv3d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1)), + activation] self.model = nn.Sequential(*model) @@ -38,32 +59,31 @@ class Encoder3d(nn.Module): return self.model(input) class BVDNet(nn.Module): - def __init__(self, N, n_downsampling=3, n_blocks=1, input_nc=3, output_nc=3,norm='batch',activation=nn.LeakyReLU(0.2)): + def __init__(self, N, n_downsampling=3, n_blocks=1, input_nc=3, output_nc=3,activation=nn.LeakyReLU(0.2)): super(BVDNet, self).__init__() ngf = 64 padding_type = 'reflect' - norm_layer = get_norm_layer(norm,'2d') - norm_layer_3d = get_norm_layer(norm,'3d') self.N = N # encoder - self.encoder3d = Encoder3d(input_nc,64,n_downsampling,norm_layer_3d,activation) - self.encoder2d = Encoder2d(input_nc,64,n_downsampling,norm_layer,activation) + self.encoder3d = Encoder3d(input_nc,64,n_downsampling,activation) + self.encoder2d = Encoder2d(input_nc,64,n_downsampling,activation) ### resnet blocks self.blocks = [] mult = 2**n_downsampling for i in range(n_blocks): - self.blocks += [ResnetBlock(ngf * mult, padding_type=padding_type, activation=nn.ReLU(True), norm_layer=norm_layer)] + self.blocks += [ResnetBlockSpectralNorm(ngf * mult, padding_type=padding_type, activation=activation)] self.blocks = nn.Sequential(*self.blocks) ### decoder self.decoder = [] for i in range(n_downsampling): mult = 2**(n_downsampling - i) - self.decoder += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1), - norm_layer(int(ngf * mult / 2)), activation] + self.decoder += [UpBlock(ngf * mult, int(ngf * mult / 2))] + # self.decoder += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1), + # norm_layer(int(ngf * mult / 2)), activation] # self.decoder += [ nn.Upsample(scale_factor = 2, mode='nearest'), # nn.ReflectionPad2d(1), # nn.Conv2d(ngf * mult, int(ngf * mult / 2),kernel_size=3, stride=1, padding=0), diff --git a/models/model_util.py b/models/model_util.py index 8865d74b74e778ee8c4ece8dd1e650daeabc6df0..f2905c7c76037c9a0c91613f6a3ac227daa8937a 100644 --- a/models/model_util.py +++ b/models/model_util.py @@ -1,6 +1,7 @@ import torch import torch.nn as nn from torch.nn import init +import torch.nn.utils.spectral_norm as SpectralNorm import functools @@ -51,4 +52,43 @@ def init_weights(net, init_type='normal', gain=0.02): init.constant_(m.bias.data, 0.0) print('initialize network with %s' % init_type) - net.apply(init_func) \ No newline at end of file + net.apply(init_func) + +class ResnetBlockSpectralNorm(nn.Module): + def __init__(self, dim, padding_type, activation=nn.LeakyReLU(0.2), use_dropout=False): + super(ResnetBlockSpectralNorm, self).__init__() + self.conv_block = self.build_conv_block(dim, padding_type, activation, use_dropout) + + def build_conv_block(self, dim, padding_type, activation, use_dropout): + conv_block = [] + p = 0 + if padding_type == 'reflect': + conv_block += [nn.ReflectionPad2d(1)] + elif padding_type == 'replicate': + conv_block += [nn.ReplicationPad2d(1)] + elif padding_type == 'zero': + p = 1 + else: + raise NotImplementedError('padding [%s] is not implemented' % padding_type) + + conv_block += [SpectralNorm(nn.Conv2d(dim, dim, kernel_size=3, padding=p)), + activation] + if use_dropout: + conv_block += [nn.Dropout(0.5)] + + p = 0 + if padding_type == 'reflect': + conv_block += [nn.ReflectionPad2d(1)] + elif padding_type == 'replicate': + conv_block += [nn.ReplicationPad2d(1)] + elif padding_type == 'zero': + p = 1 + else: + raise NotImplementedError('padding [%s] is not implemented' % padding_type) + conv_block += [SpectralNorm(nn.Conv2d(dim, dim, kernel_size=3, padding=p))] + + return nn.Sequential(*conv_block) + + def forward(self, x): + out = x + self.conv_block(x) + return out \ No newline at end of file