| | |
| | |
| |
|
| | 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)] |
| | |
| | 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): |
| | 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) |
| | |
| |
|
| | 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)] |
| | |
| | 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 |
