def set_mask(self, mask_global, opt):
print('1 in InnerCos')
mask = util.cal_feat_mask(mask_global, 3, opt.threshold)
self.mask = mask.squeeze()
if torch.cuda.is_available:
self.mask = self.mask.float().cuda()
def set_mask(self, mask_global, layer_to_last, threshold):
mask = util.cal_feat_mask(mask_global, layer_to_last, threshold)
self.mask = mask.squeeze()
return self.mask
def forward(self, input, mask):
mask = util.cal_feat_mask(mask, 3)
# input[2]:256 32 32
b, c, h, w = input[2].size()
mask_1 = torch.add(torch.neg(mask.float()), 1)
mask_1 = mask_1.expand(b, c, h, w)
x_1 = self.activation(input[0])
x_2 = self.activation(input[1])
x_3 = self.activation(input[2])
x_4 = self.activation(input[3])
x_5 = self.activation(input[4])
x_6 = self.activation(input[5])
# Change the shape of each layer and intergrate low-level/high-level features
x_1 = self.down_128(x_1)
x_2 = self.down_64(x_2)
x_3 = self.down_32(x_3)
x_4 = self.up(x_4, (32, 32))
x_5 = self.up(x_5, (32, 32))
x_6 = self.up(x_6, (32, 32))
# The first three layers are Texture/detail
# The last three layers are Structure
x_DE = torch.cat([x_1, x_2, x_3], 1)
x_ST = torch.cat([x_4, x_5, x_6], 1)
x_ST = self.down(x_ST)
x_DE = self.down(x_DE)
x_ST = [x_ST, mask_1]
x_DE = [x_DE, mask_1]
# Multi Scale PConv fill the Details
x_DE_3 = self.cov_3(x_DE)
x_DE_5 = self.cov_5(x_DE)
x_DE_7 = self.cov_7(x_DE)
x_DE_fuse = torch.cat([x_DE_3[0], x_DE_5[0], x_DE_7[0]], 1)
x_DE_fi = self.down(x_DE_fuse)
# Multi Scale PConv fill the Structure
x_ST_3 = self.cov_3(x_ST)
x_ST_5 = self.cov_5(x_ST)
x_ST_7 = self.cov_7(x_ST)
x_ST_fuse = torch.cat([x_ST_3[0], x_ST_5[0], x_ST_7[0]], 1)
x_ST_fi = self.down(x_ST_fuse)
x_cat = torch.cat([x_ST_fi, x_DE_fi], 1)
x_cat_fuse = self.fuse(x_cat)
# Feature equalizations
x_final = self.base(x_cat_fuse)
# Add back to the input
x_ST = x_final
x_DE = x_final
x_1 = self.up_128(x_DE, (128, 128)) + input[0]
x_2 = self.up_64(x_DE, (64, 64)) + input[1]
x_3 = self.up_32(x_DE, (32, 32)) + input[2]
x_4 = self.down_16(x_ST) + input[3]
x_5 = self.down_8(x_ST) + input[4]
x_6 = self.down_4(x_ST) + input[5]
out = [x_1, x_2, x_3, x_4, x_5, x_6]
loss = [x_ST_fi, x_DE_fi]
out_final = [out, loss]
return out_final
def set_mask(self, mask_global, threshold):
mask = util.cal_feat_mask(mask_global, 3, threshold)
self.mask = mask.squeeze()
if torch.cuda.is_available:
self.mask = self.mask.float().cuda()
self.mask = Variable(self.mask, requires_grad=False)
def set_mask(self, mask_global):
mask = util.cal_feat_mask(mask_global, self.layer_to_last)
self.mask = mask.squeeze()
if torch.cuda.is_available:
self.mask = self.mask.float().cuda()
def set_mask(self, mask_global):
mask = util.cal_feat_mask(mask_global, self.layer_to_last)
self.mask = mask
return self.mask
def set_mask(self, mask_global):
mask = util.cal_feat_mask(mask_global, self.layer_to_last)
self.mask = mask # 1*1*H*W (DO NOT Squeeze here!)
return self.mask
def set_mask(self, mask_global):
mask = util.cal_feat_mask(mask_global, self.layer_to_last)
self.mask = mask
self.mask_flip = torch.flip(self.mask, [3])
return self.mask
def set_mask(self, mask_global):
mask_all = util.cal_feat_mask(mask_global, self.layer_to_last)
self.mask_all = mask_all.float()
def set_mask(self, mask_global):
mask = util.cal_feat_mask(mask_global, self.layer_to_last)
self.mask = mask.squeeze()
self.mask = self.mask.float()
