Spaces:
Build error
Build error
| # Copyright (c) Microsoft Corporation. | |
| # Licensed under the MIT License. | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import os | |
| import functools | |
| from torch.autograd import Variable | |
| from util.image_pool import ImagePool | |
| from .base_model import BaseModel | |
| from . import networks | |
| import math | |
| class Mapping_Model_with_mask(nn.Module): | |
| def __init__(self, nc, mc=64, n_blocks=3, norm="instance", padding_type="reflect", opt=None): | |
| super(Mapping_Model_with_mask, self).__init__() | |
| norm_layer = networks.get_norm_layer(norm_type=norm) | |
| activation = nn.ReLU(True) | |
| model = [] | |
| tmp_nc = 64 | |
| n_up = 4 | |
| for i in range(n_up): | |
| ic = min(tmp_nc * (2 ** i), mc) | |
| oc = min(tmp_nc * (2 ** (i + 1)), mc) | |
| model += [nn.Conv2d(ic, oc, 3, 1, 1), norm_layer(oc), activation] | |
| self.before_NL = nn.Sequential(*model) | |
| if opt.NL_res: | |
| self.NL = networks.NonLocalBlock2D_with_mask_Res( | |
| mc, | |
| mc, | |
| opt.NL_fusion_method, | |
| opt.correlation_renormalize, | |
| opt.softmax_temperature, | |
| opt.use_self, | |
| opt.cosin_similarity, | |
| ) | |
| print("You are using NL + Res") | |
| model = [] | |
| for i in range(n_blocks): | |
| model += [ | |
| networks.ResnetBlock( | |
| mc, | |
| padding_type=padding_type, | |
| activation=activation, | |
| norm_layer=norm_layer, | |
| opt=opt, | |
| dilation=opt.mapping_net_dilation, | |
| ) | |
| ] | |
| for i in range(n_up - 1): | |
| ic = min(64 * (2 ** (4 - i)), mc) | |
| oc = min(64 * (2 ** (3 - i)), mc) | |
| model += [nn.Conv2d(ic, oc, 3, 1, 1), norm_layer(oc), activation] | |
| model += [nn.Conv2d(tmp_nc * 2, tmp_nc, 3, 1, 1)] | |
| if opt.feat_dim > 0 and opt.feat_dim < 64: | |
| model += [norm_layer(tmp_nc), activation, nn.Conv2d(tmp_nc, opt.feat_dim, 1, 1)] | |
| # model += [nn.Conv2d(64, 1, 1, 1, 0)] | |
| self.after_NL = nn.Sequential(*model) | |
| def forward(self, input, mask): | |
| x1 = self.before_NL(input) | |
| del input | |
| x2 = self.NL(x1, mask) | |
| del x1, mask | |
| x3 = self.after_NL(x2) | |
| del x2 | |
| return x3 | |
| class Mapping_Model_with_mask_2(nn.Module): ## Multi-Scale Patch Attention | |
| def __init__(self, nc, mc=64, n_blocks=3, norm="instance", padding_type="reflect", opt=None): | |
| super(Mapping_Model_with_mask_2, self).__init__() | |
| norm_layer = networks.get_norm_layer(norm_type=norm) | |
| activation = nn.ReLU(True) | |
| model = [] | |
| tmp_nc = 64 | |
| n_up = 4 | |
| for i in range(n_up): | |
| ic = min(tmp_nc * (2 ** i), mc) | |
| oc = min(tmp_nc * (2 ** (i + 1)), mc) | |
| model += [nn.Conv2d(ic, oc, 3, 1, 1), norm_layer(oc), activation] | |
| for i in range(2): | |
| model += [ | |
| networks.ResnetBlock( | |
| mc, | |
| padding_type=padding_type, | |
| activation=activation, | |
| norm_layer=norm_layer, | |
| opt=opt, | |
| dilation=opt.mapping_net_dilation, | |
| ) | |
| ] | |
| print("Mapping: You are using multi-scale patch attention, conv combine + mask input") | |
| self.before_NL = nn.Sequential(*model) | |
| if opt.mapping_exp==1: | |
| self.NL_scale_1=networks.Patch_Attention_4(mc,mc,8) | |
| model = [] | |
| for i in range(2): | |
| model += [ | |
| networks.ResnetBlock( | |
| mc, | |
| padding_type=padding_type, | |
| activation=activation, | |
| norm_layer=norm_layer, | |
| opt=opt, | |
| dilation=opt.mapping_net_dilation, | |
| ) | |
| ] | |
| self.res_block_1 = nn.Sequential(*model) | |
| if opt.mapping_exp==1: | |
| self.NL_scale_2=networks.Patch_Attention_4(mc,mc,4) | |
| model = [] | |
| for i in range(2): | |
| model += [ | |
| networks.ResnetBlock( | |
| mc, | |
| padding_type=padding_type, | |
| activation=activation, | |
| norm_layer=norm_layer, | |
| opt=opt, | |
| dilation=opt.mapping_net_dilation, | |
| ) | |
| ] | |
| self.res_block_2 = nn.Sequential(*model) | |
| if opt.mapping_exp==1: | |
| self.NL_scale_3=networks.Patch_Attention_4(mc,mc,2) | |
| # self.NL_scale_3=networks.Patch_Attention_2(mc,mc,2) | |
| model = [] | |
| for i in range(2): | |
| model += [ | |
| networks.ResnetBlock( | |
| mc, | |
| padding_type=padding_type, | |
| activation=activation, | |
| norm_layer=norm_layer, | |
| opt=opt, | |
| dilation=opt.mapping_net_dilation, | |
| ) | |
| ] | |
| for i in range(n_up - 1): | |
| ic = min(64 * (2 ** (4 - i)), mc) | |
| oc = min(64 * (2 ** (3 - i)), mc) | |
| model += [nn.Conv2d(ic, oc, 3, 1, 1), norm_layer(oc), activation] | |
| model += [nn.Conv2d(tmp_nc * 2, tmp_nc, 3, 1, 1)] | |
| if opt.feat_dim > 0 and opt.feat_dim < 64: | |
| model += [norm_layer(tmp_nc), activation, nn.Conv2d(tmp_nc, opt.feat_dim, 1, 1)] | |
| # model += [nn.Conv2d(64, 1, 1, 1, 0)] | |
| self.after_NL = nn.Sequential(*model) | |
| def forward(self, input, mask): | |
| x1 = self.before_NL(input) | |
| x2 = self.NL_scale_1(x1,mask) | |
| x3 = self.res_block_1(x2) | |
| x4 = self.NL_scale_2(x3,mask) | |
| x5 = self.res_block_2(x4) | |
| x6 = self.NL_scale_3(x5,mask) | |
| x7 = self.after_NL(x6) | |
| return x7 | |
| def inference_forward(self, input, mask): | |
| x1 = self.before_NL(input) | |
| del input | |
| x2 = self.NL_scale_1.inference_forward(x1,mask) | |
| del x1 | |
| x3 = self.res_block_1(x2) | |
| del x2 | |
| x4 = self.NL_scale_2.inference_forward(x3,mask) | |
| del x3 | |
| x5 = self.res_block_2(x4) | |
| del x4 | |
| x6 = self.NL_scale_3.inference_forward(x5,mask) | |
| del x5 | |
| x7 = self.after_NL(x6) | |
| del x6 | |
| return x7 |