Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torch.nn.parallel | |
| from torch.autograd import Variable | |
| import torch.nn.functional as F | |
| from torchvision import models | |
| import torch.utils.model_zoo as model_zoo | |
| from torch.nn import init | |
| import os | |
| import numpy as np | |
| def weights_init_normal(m): | |
| classname = m.__class__.__name__ | |
| if classname.find('Conv') != -1: | |
| init.normal_(m.weight.data, 0.0, 0.02) | |
| elif classname.find('Linear') != -1: | |
| init.normal(m.weight.data, 0.0, 0.02) | |
| elif classname.find('BatchNorm2d') != -1: | |
| init.normal_(m.weight.data, 1.0, 0.02) | |
| init.constant_(m.bias.data, 0.0) | |
| def weights_init_xavier(m): | |
| classname = m.__class__.__name__ | |
| if classname.find('Conv') != -1: | |
| init.xavier_normal_(m.weight.data, gain=0.02) | |
| elif classname.find('Linear') != -1: | |
| init.xavier_normal_(m.weight.data, gain=0.02) | |
| elif classname.find('BatchNorm2d') != -1: | |
| init.normal_(m.weight.data, 1.0, 0.02) | |
| init.constant_(m.bias.data, 0.0) | |
| def weights_init_kaiming(m): | |
| classname = m.__class__.__name__ | |
| if classname.find('Conv') != -1: | |
| init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') | |
| elif classname.find('Linear') != -1: | |
| init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') | |
| elif classname.find('BatchNorm2d') != -1: | |
| init.normal_(m.weight.data, 1.0, 0.02) | |
| init.constant_(m.bias.data, 0.0) | |
| def init_weights(net, init_type='normal'): | |
| print('initialization method [%s]' % init_type) | |
| if init_type == 'normal': | |
| net.apply(weights_init_normal) | |
| elif init_type == 'xavier': | |
| net.apply(weights_init_xavier) | |
| elif init_type == 'kaiming': | |
| net.apply(weights_init_kaiming) | |
| else: | |
| raise NotImplementedError('initialization method [%s] is not implemented' % init_type) | |
| class FeatureExtraction(nn.Module): | |
| def __init__(self, input_nc, ngf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_dropout=False): | |
| super(FeatureExtraction, self).__init__() | |
| downconv = nn.Conv2d(input_nc, ngf, kernel_size=4, stride=2, padding=1) | |
| model = [downconv, nn.ReLU(True), norm_layer(ngf)] | |
| for i in range(n_layers): | |
| in_ngf = 2 ** i * ngf if 2 ** i * ngf < 512 else 512 | |
| out_ngf = 2 ** (i + 1) * ngf if 2 ** i * ngf < 512 else 512 | |
| downconv = nn.Conv2d(in_ngf, out_ngf, kernel_size=4, stride=2, padding=1) | |
| model += [downconv, nn.ReLU(True)] | |
| model += [norm_layer(out_ngf)] | |
| model += [nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1), nn.ReLU(True)] | |
| model += [norm_layer(512)] | |
| model += [nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1), nn.ReLU(True)] | |
| self.model = nn.Sequential(*model) | |
| init_weights(self.model, init_type='normal') | |
| def forward(self, x): | |
| return self.model(x) | |
| class FeatureL2Norm(torch.nn.Module): | |
| def __init__(self): | |
| super(FeatureL2Norm, self).__init__() | |
| def forward(self, feature): | |
| epsilon = 1e-6 | |
| norm = torch.pow(torch.sum(torch.pow(feature, 2), 1) + epsilon, 0.5).unsqueeze(1).expand_as(feature) | |
| return torch.div(feature, norm) | |
| class FeatureCorrelation(nn.Module): | |
| def __init__(self): | |
| super(FeatureCorrelation, self).__init__() | |
| def forward(self, feature_A, feature_B): | |
| b, c, h, w = feature_A.size() | |
| # reshape features for matrix multiplication | |
| feature_A = feature_A.transpose(2, 3).contiguous().view(b, c, h * w) | |
| feature_B = feature_B.view(b, c, h * w).transpose(1, 2) | |
| # perform matrix mult. | |
| feature_mul = torch.bmm(feature_B, feature_A) | |
| correlation_tensor = feature_mul.view(b, h, w, h * w).transpose(2, 3).transpose(1, 2) | |
| return correlation_tensor | |
| class FeatureRegression(nn.Module): | |
| def __init__(self, input_nc=512, output_dim=6, use_cuda=True): | |
| super(FeatureRegression, self).__init__() | |
| self.conv = nn.Sequential( | |
| nn.Conv2d(input_nc, 512, kernel_size=4, stride=2, padding=1), | |
| nn.BatchNorm2d(512), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(512, 256, kernel_size=4, stride=2, padding=1), | |
| nn.BatchNorm2d(256), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(256, 128, kernel_size=3, padding=1), | |
| nn.BatchNorm2d(128), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(128, 64, kernel_size=3, padding=1), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True), | |
| ) | |
| self.linear = nn.Linear(64 * 4 * 3, output_dim) | |
| self.tanh = nn.Tanh() | |
| if use_cuda: | |
| self.conv.cuda() | |
| self.linear.cuda() | |
| self.tanh.cuda() | |
| def forward(self, x): | |
| x = self.conv(x) | |
| x = x.view(x.size(0), -1) | |
| x = self.linear(x) | |
| x = self.tanh(x) | |
| return x | |
| class AffineGridGen(nn.Module): | |
| def __init__(self, out_h=256, out_w=192, out_ch=3): | |
| super(AffineGridGen, self).__init__() | |
| self.out_h = out_h | |
| self.out_w = out_w | |
| self.out_ch = out_ch | |
| def forward(self, theta): | |
| theta = theta.contiguous() | |
| batch_size = theta.size()[0] | |
| out_size = torch.Size((batch_size, self.out_ch, self.out_h, self.out_w)) | |
| return F.affine_grid(theta, out_size) | |
| class TpsGridGen(nn.Module): | |
| def __init__(self, out_h=256, out_w=192, use_regular_grid=True, grid_size=3, reg_factor=0, use_cuda=True): | |
| super(TpsGridGen, self).__init__() | |
| self.out_h, self.out_w = out_h, out_w | |
| self.reg_factor = reg_factor | |
| self.use_cuda = use_cuda | |
| # create grid in numpy | |
| self.grid = np.zeros([self.out_h, self.out_w, 3], dtype=np.float32) | |
| # sampling grid with dim-0 coords (Y) | |
| self.grid_X, self.grid_Y = np.meshgrid(np.linspace(-1, 1, out_w), np.linspace(-1, 1, out_h)) | |
| # grid_X,grid_Y: size [1,H,W,1,1] | |
| self.grid_X = torch.FloatTensor(self.grid_X).unsqueeze(0).unsqueeze(3) | |
| self.grid_Y = torch.FloatTensor(self.grid_Y).unsqueeze(0).unsqueeze(3) | |
| if use_cuda: | |
| self.grid_X = self.grid_X.cuda() | |
| self.grid_Y = self.grid_Y.cuda() | |
| # initialize regular grid for control points P_i | |
| if use_regular_grid: | |
| axis_coords = np.linspace(-1, 1, grid_size) | |
| self.N = grid_size * grid_size | |
| P_Y, P_X = np.meshgrid(axis_coords, axis_coords) | |
| P_X = np.reshape(P_X, (-1, 1)) # size (N,1) | |
| P_Y = np.reshape(P_Y, (-1, 1)) # size (N,1) | |
| P_X = torch.FloatTensor(P_X) | |
| P_Y = torch.FloatTensor(P_Y) | |
| self.P_X_base = P_X.clone() | |
| self.P_Y_base = P_Y.clone() | |
| self.Li = self.compute_L_inverse(P_X, P_Y).unsqueeze(0) | |
| self.P_X = P_X.unsqueeze(2).unsqueeze(3).unsqueeze(4).transpose(0, 4) | |
| self.P_Y = P_Y.unsqueeze(2).unsqueeze(3).unsqueeze(4).transpose(0, 4) | |
| if use_cuda: | |
| self.P_X = self.P_X.cuda() | |
| self.P_Y = self.P_Y.cuda() | |
| self.P_X_base = self.P_X_base.cuda() | |
| self.P_Y_base = self.P_Y_base.cuda() | |
| def forward(self, theta): | |
| warped_grid = self.apply_transformation(theta, torch.cat((self.grid_X, self.grid_Y), 3)) | |
| return warped_grid | |
| def compute_L_inverse(self, X, Y): | |
| N = X.size()[0] # num of points (along dim 0) | |
| # construct matrix K | |
| Xmat = X.expand(N, N) | |
| Ymat = Y.expand(N, N) | |
| P_dist_squared = torch.pow(Xmat - Xmat.transpose(0, 1), 2) + torch.pow(Ymat - Ymat.transpose(0, 1), 2) | |
| P_dist_squared[P_dist_squared == 0] = 1 # make diagonal 1 to avoid NaN in log computation | |
| K = torch.mul(P_dist_squared, torch.log(P_dist_squared)) | |
| # construct matrix L | |
| O = torch.FloatTensor(N, 1).fill_(1) | |
| Z = torch.FloatTensor(3, 3).fill_(0) | |
| P = torch.cat((O, X, Y), 1) | |
| L = torch.cat((torch.cat((K, P), 1), torch.cat((P.transpose(0, 1), Z), 1)), 0) | |
| Li = torch.inverse(L) | |
| if self.use_cuda: | |
| Li = Li.cuda() | |
| return Li | |
| def apply_transformation(self, theta, points): | |
| if theta.dim() == 2: | |
| theta = theta.unsqueeze(2).unsqueeze(3) | |
| # points should be in the [B,H,W,2] format, | |
| # where points[:,:,:,0] are the X coords | |
| # and points[:,:,:,1] are the Y coords | |
| # input are the corresponding control points P_i | |
| batch_size = theta.size()[0] | |
| # split theta into point coordinates | |
| Q_X = theta[:, :self.N, :, :].squeeze(3) | |
| Q_Y = theta[:, self.N:, :, :].squeeze(3) | |
| Q_X = Q_X + self.P_X_base.expand_as(Q_X) | |
| Q_Y = Q_Y + self.P_Y_base.expand_as(Q_Y) | |
| # get spatial dimensions of points | |
| points_b = points.size()[0] | |
| points_h = points.size()[1] | |
| points_w = points.size()[2] | |
| # repeat pre-defined control points along spatial dimensions of points to be transformed | |
| P_X = self.P_X.expand((1, points_h, points_w, 1, self.N)) | |
| P_Y = self.P_Y.expand((1, points_h, points_w, 1, self.N)) | |
| # compute weigths for non-linear part | |
| W_X = torch.bmm(self.Li[:, :self.N, :self.N].expand((batch_size, self.N, self.N)), Q_X) | |
| W_Y = torch.bmm(self.Li[:, :self.N, :self.N].expand((batch_size, self.N, self.N)), Q_Y) | |
| # reshape | |
| # W_X,W,Y: size [B,H,W,1,N] | |
| W_X = W_X.unsqueeze(3).unsqueeze(4).transpose(1, 4).repeat(1, points_h, points_w, 1, 1) | |
| W_Y = W_Y.unsqueeze(3).unsqueeze(4).transpose(1, 4).repeat(1, points_h, points_w, 1, 1) | |
| # compute weights for affine part | |
| A_X = torch.bmm(self.Li[:, self.N:, :self.N].expand((batch_size, 3, self.N)), Q_X) | |
| A_Y = torch.bmm(self.Li[:, self.N:, :self.N].expand((batch_size, 3, self.N)), Q_Y) | |
| # reshape | |
| # A_X,A,Y: size [B,H,W,1,3] | |
| A_X = A_X.unsqueeze(3).unsqueeze(4).transpose(1, 4).repeat(1, points_h, points_w, 1, 1) | |
| A_Y = A_Y.unsqueeze(3).unsqueeze(4).transpose(1, 4).repeat(1, points_h, points_w, 1, 1) | |
| # compute distance P_i - (grid_X,grid_Y) | |
| # grid is expanded in point dim 4, but not in batch dim 0, as points P_X,P_Y are fixed for all batch | |
| points_X_for_summation = points[:, :, :, 0].unsqueeze(3).unsqueeze(4).expand( | |
| points[:, :, :, 0].size() + (1, self.N)) | |
| points_Y_for_summation = points[:, :, :, 1].unsqueeze(3).unsqueeze(4).expand( | |
| points[:, :, :, 1].size() + (1, self.N)) | |
| if points_b == 1: | |
| delta_X = points_X_for_summation - P_X | |
| delta_Y = points_Y_for_summation - P_Y | |
| else: | |
| # use expanded P_X,P_Y in batch dimension | |
| delta_X = points_X_for_summation - P_X.expand_as(points_X_for_summation) | |
| delta_Y = points_Y_for_summation - P_Y.expand_as(points_Y_for_summation) | |
| dist_squared = torch.pow(delta_X, 2) + torch.pow(delta_Y, 2) | |
| # U: size [1,H,W,1,N] | |
| dist_squared[dist_squared == 0] = 1 # avoid NaN in log computation | |
| U = torch.mul(dist_squared, torch.log(dist_squared)) | |
| # expand grid in batch dimension if necessary | |
| points_X_batch = points[:, :, :, 0].unsqueeze(3) | |
| points_Y_batch = points[:, :, :, 1].unsqueeze(3) | |
| if points_b == 1: | |
| points_X_batch = points_X_batch.expand((batch_size,) + points_X_batch.size()[1:]) | |
| points_Y_batch = points_Y_batch.expand((batch_size,) + points_Y_batch.size()[1:]) | |
| points_X_prime = A_X[:, :, :, :, 0] + \ | |
| torch.mul(A_X[:, :, :, :, 1], points_X_batch) + \ | |
| torch.mul(A_X[:, :, :, :, 2], points_Y_batch) + \ | |
| torch.sum(torch.mul(W_X, U.expand_as(W_X)), 4) | |
| points_Y_prime = A_Y[:, :, :, :, 0] + \ | |
| torch.mul(A_Y[:, :, :, :, 1], points_X_batch) + \ | |
| torch.mul(A_Y[:, :, :, :, 2], points_Y_batch) + \ | |
| torch.sum(torch.mul(W_Y, U.expand_as(W_Y)), 4) | |
| return torch.cat((points_X_prime, points_Y_prime), 3) | |
| # Defines the Unet generator. | |
| # |num_downs|: number of downsamplings in UNet. For example, | |
| # if |num_downs| == 7, image of size 128x128 will become of size 1x1 | |
| # at the bottleneck | |
| class UnetGenerator(nn.Module): | |
| def __init__(self, input_nc, output_nc, num_downs, ngf=64, | |
| norm_layer=nn.BatchNorm2d, use_dropout=False): | |
| super(UnetGenerator, self).__init__() | |
| # construct unet structure | |
| unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, | |
| innermost=True) | |
| for i in range(num_downs - 5): | |
| unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, | |
| norm_layer=norm_layer, use_dropout=use_dropout) | |
| unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, | |
| norm_layer=norm_layer) | |
| unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, | |
| norm_layer=norm_layer) | |
| unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer) | |
| unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, | |
| norm_layer=norm_layer) | |
| self.model = unet_block | |
| def forward(self, input): | |
| return self.model(input) | |
| # Defines the submodule with skip connection. | |
| # X -------------------identity---------------------- X | |
| # |-- downsampling -- |submodule| -- upsampling --| | |
| class UnetSkipConnectionBlock(nn.Module): | |
| def __init__(self, outer_nc, inner_nc, input_nc=None, | |
| submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False): | |
| super(UnetSkipConnectionBlock, self).__init__() | |
| self.outermost = outermost | |
| use_bias = norm_layer == nn.InstanceNorm2d | |
| if input_nc is None: | |
| input_nc = outer_nc | |
| downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4, | |
| stride=2, padding=1, bias=use_bias) | |
| downrelu = nn.LeakyReLU(0.2, True) | |
| uprelu = nn.ReLU(True) | |
| if norm_layer != None: | |
| downnorm = norm_layer(inner_nc) | |
| upnorm = norm_layer(outer_nc) | |
| if outermost: | |
| upsample = nn.Upsample(scale_factor=2, mode='bilinear') | |
| upconv = nn.Conv2d(inner_nc * 2, outer_nc, kernel_size=3, stride=1, padding=1, bias=use_bias) | |
| down = [downconv] | |
| # up = [uprelu, upsample, upconv, upnorm] | |
| up = [uprelu, upsample, upconv] | |
| model = down + [submodule] + up | |
| elif innermost: | |
| upsample = nn.Upsample(scale_factor=2, mode='bilinear') | |
| upconv = nn.Conv2d(inner_nc, outer_nc, kernel_size=3, stride=1, padding=1, bias=use_bias) | |
| down = [downrelu, downconv] | |
| if norm_layer == None: | |
| up = [uprelu, upsample, upconv] | |
| else: | |
| up = [uprelu, upsample, upconv, upnorm] | |
| model = down + up | |
| else: | |
| upsample = nn.Upsample(scale_factor=2, mode='bilinear') | |
| upconv = nn.Conv2d(inner_nc * 2, outer_nc, kernel_size=3, stride=1, padding=1, bias=use_bias) | |
| if norm_layer == None: | |
| down = [downrelu, downconv] | |
| up = [uprelu, upsample, upconv] | |
| else: | |
| down = [downrelu, downconv, downnorm] | |
| up = [uprelu, upsample, upconv, upnorm] | |
| if use_dropout: | |
| model = down + [submodule] + up + [nn.Dropout(0.5)] | |
| else: | |
| model = down + [submodule] + up | |
| self.model = nn.Sequential(*model) | |
| def forward(self, x): | |
| if self.outermost: | |
| return self.model(x) | |
| else: | |
| return torch.cat([x, self.model(x)], 1) | |
| # UNet with residual blocks | |
| class ResidualBlock(nn.Module): | |
| def __init__(self, in_features=64, norm_layer=nn.BatchNorm2d): | |
| super(ResidualBlock, self).__init__() | |
| self.relu = nn.ReLU(True) | |
| if norm_layer == None: | |
| # hard to converge with out batch or instance norm | |
| self.block = nn.Sequential( | |
| nn.Conv2d(in_features, in_features, 3, 1, 1, bias=False), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(in_features, in_features, 3, 1, 1, bias=False), | |
| ) | |
| else: | |
| self.block = nn.Sequential( | |
| nn.Conv2d(in_features, in_features, 3, 1, 1, bias=False), | |
| norm_layer(in_features), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(in_features, in_features, 3, 1, 1, bias=False), | |
| norm_layer(in_features) | |
| ) | |
| def forward(self, x): | |
| residual = x | |
| out = self.block(x) | |
| out += residual | |
| out = self.relu(out) | |
| return out | |
| # return self.relu(x + self.block(x)) | |
| class ResUnetGenerator(nn.Module): | |
| def __init__(self, input_nc, output_nc, num_downs, ngf=64, | |
| norm_layer=nn.BatchNorm2d, use_dropout=False): | |
| super(ResUnetGenerator, self).__init__() | |
| # construct unet structure | |
| unet_block = ResUnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, | |
| innermost=True) | |
| for i in range(num_downs - 5): | |
| unet_block = ResUnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, | |
| norm_layer=norm_layer, use_dropout=use_dropout) | |
| unet_block = ResUnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, | |
| norm_layer=norm_layer) | |
| unet_block = ResUnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, | |
| norm_layer=norm_layer) | |
| unet_block = ResUnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, | |
| norm_layer=norm_layer) | |
| unet_block = ResUnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, | |
| norm_layer=norm_layer) | |
| self.model = unet_block | |
| def forward(self, input): | |
| output = self.model(input) | |
| # print("\tIn Model: input size", input.size(), | |
| # "output size", output.size()) | |
| return output | |
| # Defines the submodule with skip connection. | |
| # X -------------------identity---------------------- X | |
| # |-- downsampling -- |submodule| -- upsampling --| | |
| class ResUnetSkipConnectionBlock(nn.Module): | |
| def __init__(self, outer_nc, inner_nc, input_nc=None, | |
| submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False): | |
| super(ResUnetSkipConnectionBlock, self).__init__() | |
| self.outermost = outermost | |
| use_bias = norm_layer == nn.InstanceNorm2d | |
| if input_nc is None: | |
| input_nc = outer_nc | |
| downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=3, | |
| stride=2, padding=1, bias=use_bias) | |
| # add two resblock | |
| res_downconv = [ResidualBlock(inner_nc, norm_layer), ResidualBlock(inner_nc, norm_layer)] | |
| res_upconv = [ResidualBlock(outer_nc, norm_layer), ResidualBlock(outer_nc, norm_layer)] | |
| # res_downconv = [ResidualBlock(inner_nc)] | |
| # res_upconv = [ResidualBlock(outer_nc)] | |
| downrelu = nn.ReLU(True) | |
| uprelu = nn.ReLU(True) | |
| if norm_layer != None: | |
| downnorm = norm_layer(inner_nc) | |
| upnorm = norm_layer(outer_nc) | |
| if outermost: | |
| upsample = nn.Upsample(scale_factor=2, mode='nearest') | |
| upconv = nn.Conv2d(inner_nc * 2, outer_nc, kernel_size=3, stride=1, padding=1, bias=use_bias) | |
| down = [downconv, downrelu] + res_downconv | |
| # up = [uprelu, upsample, upconv, upnorm] | |
| up = [upsample, upconv] | |
| model = down + [submodule] + up | |
| elif innermost: | |
| upsample = nn.Upsample(scale_factor=2, mode='nearest') | |
| upconv = nn.Conv2d(inner_nc, outer_nc, kernel_size=3, stride=1, padding=1, bias=use_bias) | |
| down = [downconv, downrelu] + res_downconv | |
| if norm_layer == None: | |
| up = [upsample, upconv, uprelu] + res_upconv | |
| else: | |
| up = [upsample, upconv, upnorm, uprelu] + res_upconv | |
| model = down + up | |
| else: | |
| upsample = nn.Upsample(scale_factor=2, mode='nearest') | |
| upconv = nn.Conv2d(inner_nc * 2, outer_nc, kernel_size=3, stride=1, padding=1, bias=use_bias) | |
| if norm_layer == None: | |
| down = [downconv, downrelu] + res_downconv | |
| up = [upsample, upconv, uprelu] + res_upconv | |
| else: | |
| down = [downconv, downnorm, downrelu] + res_downconv | |
| up = [upsample, upconv, upnorm, uprelu] + res_upconv | |
| if use_dropout: | |
| model = down + [submodule] + up + [nn.Dropout(0.5)] | |
| else: | |
| model = down + [submodule] + up | |
| self.model = nn.Sequential(*model) | |
| def forward(self, x): | |
| if self.outermost: | |
| return self.model(x) | |
| else: | |
| return torch.cat([x, self.model(x)], 1) | |
| class Vgg19(nn.Module): | |
| def __init__(self, requires_grad=False): | |
| super(Vgg19, self).__init__() | |
| vgg_pretrained_features = models.vgg19(pretrained=True).features | |
| self.slice1 = nn.Sequential() | |
| self.slice2 = nn.Sequential() | |
| self.slice3 = nn.Sequential() | |
| self.slice4 = nn.Sequential() | |
| self.slice5 = nn.Sequential() | |
| for x in range(2): | |
| self.slice1.add_module(str(x), vgg_pretrained_features[x]) | |
| for x in range(2, 7): | |
| self.slice2.add_module(str(x), vgg_pretrained_features[x]) | |
| for x in range(7, 12): | |
| self.slice3.add_module(str(x), vgg_pretrained_features[x]) | |
| for x in range(12, 21): | |
| self.slice4.add_module(str(x), vgg_pretrained_features[x]) | |
| for x in range(21, 30): | |
| self.slice5.add_module(str(x), vgg_pretrained_features[x]) | |
| if not requires_grad: | |
| for param in self.parameters(): | |
| param.requires_grad = False | |
| def forward(self, X): | |
| h_relu1 = self.slice1(X) | |
| h_relu2 = self.slice2(h_relu1) | |
| h_relu3 = self.slice3(h_relu2) | |
| h_relu4 = self.slice4(h_relu3) | |
| h_relu5 = self.slice5(h_relu4) | |
| out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5] | |
| return out | |
| def gram_matrix(input): | |
| a, b, c, d = input.size() # a=batch size(=1) | |
| # b=number of feature maps | |
| # (c,d)=dimensions of a f. map (N=c*d) | |
| features = input.view(a * b, c * d) # resise F_XL into \hat F_XL | |
| G = torch.mm(features, features.t()) # compute the gram product | |
| # we 'normalize' the values of the gram matrix | |
| # by dividing by the number of element in each feature maps. | |
| return G.div(a * b * c * d) | |
| class StyleLoss(nn.Module): | |
| def __init__(self): | |
| super(StyleLoss, self).__init__() | |
| def forward(self, x, y): | |
| Gx = gram_matrix(x) | |
| Gy = gram_matrix(y) | |
| return F.mse_loss(Gx, Gy) * 30000000 | |
| class VGGLoss(nn.Module): | |
| def __init__(self, model=None): | |
| super(VGGLoss, self).__init__() | |
| if model is None: | |
| self.vgg = Vgg19() | |
| else: | |
| self.vgg = model | |
| self.vgg.cuda() | |
| # self.vgg.eval() | |
| self.criterion = nn.L1Loss() | |
| self.style_criterion = StyleLoss() | |
| self.weights = [1.0, 1.0, 1.0, 1.0, 1.0] | |
| self.style_weights = [1.0, 1.0, 1.0, 1.0, 1.0] | |
| # self.weights = [5.0, 1.0, 0.5, 0.4, 0.8] | |
| # self.style_weights = [10e4, 1000, 50, 15, 50] | |
| def forward(self, x, y, style=False): | |
| x_vgg, y_vgg = self.vgg(x), self.vgg(y) | |
| loss = 0 | |
| if style: | |
| # return both perceptual loss and style loss. | |
| style_loss = 0 | |
| for i in range(len(x_vgg)): | |
| this_loss = (self.weights[i] * | |
| self.criterion(x_vgg[i], y_vgg[i].detach())) | |
| this_style_loss = (self.style_weights[i] * | |
| self.style_criterion(x_vgg[i], y_vgg[i].detach())) | |
| loss += this_loss | |
| style_loss += this_style_loss | |
| return loss, style_loss | |
| for i in range(len(x_vgg)): | |
| this_loss = (self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())) | |
| loss += this_loss | |
| return loss | |
| class GMM(nn.Module): | |
| """ Geometric Matching Module | |
| """ | |
| def __init__(self, opt, input_nc): | |
| super(GMM, self).__init__() | |
| self.extractionA = FeatureExtraction(input_nc, ngf=64, n_layers=3, norm_layer=nn.BatchNorm2d) | |
| self.extractionB = FeatureExtraction(3, ngf=64, n_layers=3, norm_layer=nn.BatchNorm2d) | |
| self.l2norm = FeatureL2Norm() | |
| self.correlation = FeatureCorrelation() | |
| self.regression = FeatureRegression(input_nc=192, output_dim=2 * opt.grid_size ** 2, use_cuda=True) | |
| self.gridGen = TpsGridGen(opt.fine_height, opt.fine_width, use_cuda=True, grid_size=opt.grid_size) | |
| def forward(self, inputA, inputB): | |
| featureA = self.extractionA(inputA) | |
| featureB = self.extractionB(inputB) | |
| featureA = self.l2norm(featureA) | |
| featureB = self.l2norm(featureB) | |
| correlation = self.correlation(featureA, featureB) | |
| theta = self.regression(correlation) | |
| grid = self.gridGen(theta) | |
| return grid, theta | |
| def save_checkpoint(model, save_path): | |
| if not os.path.exists(os.path.dirname(save_path)): | |
| os.makedirs(os.path.dirname(save_path)) | |
| torch.save(model.state_dict(), save_path) | |
| def load_checkpoint(model, checkpoint_path): | |
| if not os.path.exists(checkpoint_path): | |
| print('No checkpoint!') | |
| return | |
| model.load_state_dict(torch.load(checkpoint_path)) | |
| # try: | |
| # model.load_state_dict(torch.load(checkpoint_path)) | |
| # except: | |
| # model = nn.DataParallel(model) | |
| # model.load_state_dict(torch.load(checkpoint_path)) | |