| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from .base_network import BaseNetwork |
| import random |
| from .blocks import Conv2dBlock |
|
|
| def inst_id_to_label_id(inst_id: int) -> int: |
| return int(inst_id) % 120 |
|
|
| class DeeperZencoder(BaseNetwork): |
| def __init__(self,cfg, input_nc = 3, output_nc = 512, ngf=32, n_downsampling=2, norm_layer=nn.InstanceNorm2d): |
| super().__init__() |
| self.cfg = cfg |
| self.n_downsampling = 6 |
| self.non_spade_norm_layer = norm_layer |
| self.output_nc = cfg["style_length"] |
| input_nc = input_nc +1 |
|
|
| self.gamma = nn.Linear(1, 1) |
| self.beta = nn.Linear(1, 1) |
| |
| input_nc = cfg['input_nc'] |
| ngf = cfg['ngf'] |
| output_nc = cfg['style_length'] |
| lab_nc = cfg['lab_dim'] + 1 |
| g_norm = cfg['G_norm_type'] |
| self.enc1 = Conv2dBlock(4, 16, kernel_size=3, stride=1, padding=1, norm=g_norm, activation='lrelu') |
| self.enc2 = Conv2dBlock(16, 32, kernel_size=3, stride=2, padding=1, norm=g_norm, activation='lrelu') |
| self.enc3 = Conv2dBlock(32, 64, kernel_size=3, stride=2, padding=1, norm=g_norm, activation='lrelu') |
| self.enc4 = Conv2dBlock(64, 128, kernel_size=3, stride=2, padding=1, norm=g_norm, activation='lrelu') |
| self.enc5 = Conv2dBlock(128, 256, kernel_size=3, stride=2, padding=1, norm=g_norm, activation='lrelu') |
| self.enc6 = Conv2dBlock(256, 512, kernel_size=3, stride=1, padding=1, dilation=1, norm=g_norm, activation='lrelu') |
|
|
| |
| self.dec5 = Conv2dBlock(512+128, 512, kernel_size=3, stride=1, padding=1, norm=g_norm, |
| activation='lrelu') |
| self.dec4 = Conv2dBlock(512+64, 256, kernel_size=3, stride=1, padding=1, norm=g_norm, |
| activation='lrelu') |
| self.dec3 = Conv2dBlock(256+32, 256, kernel_size=3, stride=1, padding=1, norm=g_norm, |
| activation='lrelu') |
| self.dec2 = Conv2dBlock(256+16, 256, kernel_size=3, stride=1, padding=1, norm=g_norm, |
| activation='lrelu') |
| self.dec1 = Conv2dBlock(256, output_nc, kernel_size=3, stride=1, padding=1, norm='none', activation='tanh') |
|
|
| def forward(self, input, segmap, valids, instance_map=None, cached_codes=None): |
| masked_input = input*valids |
| masked_input = torch.cat((masked_input, valids), 1) |
| codes = self._forward_layers(masked_input) |
|
|
| original_input = torch.cat((input, torch.ones(valids.shape, device=valids.device, dtype=valids.dtype)), 1) |
| unmasked_codes = self._forward_layers(original_input) |
|
|
| if instance_map is None: |
| styles_map = segmap |
| else: |
| styles_map = instance_map |
| styles_map = F.interpolate(styles_map, size=codes.size()[2:], mode='nearest') |
| valids = F.interpolate(valids, size=codes.size()[2:], mode='nearest') |
| |
| style_codes = codes.clone() |
| for b in range(input.size()[0]): |
| inst_list = torch.unique(styles_map[b]).to(torch.long) |
| for i in inst_list: |
| indices = (styles_map[b:b+1] == int(i)).nonzero() |
| valid_ins = valids[indices[:,0] + b, :, indices[:,2], indices[:,3]] |
| if valid_ins.sum() > 0: |
| output_ins = codes[indices[:,0] + b, :, indices[:,2], indices[:,3]] |
| mean_feat = output_ins.mean(dim=0).expand_as(output_ins) |
| valid_ratio = valid_ins.mean().unsqueeze(0) |
| mean_feat = mean_feat * torch.sigmoid(self.gamma(valid_ratio)) + torch.sigmoid(self.beta(valid_ratio)) |
| elif self.cfg["is_train"]: |
| unmasked_output_ins = unmasked_codes[indices[:,0] + b, :, indices[:,2], indices[:,3]] |
| mean_feat = unmasked_output_ins.mean(dim=0).expand_as(unmasked_output_ins) |
| valid_ratio = torch.ones((1,)) |
| mean_feat = mean_feat * torch.sigmoid(self.gamma(valid_ratio)) + torch.sigmoid(self.beta(valid_ratio)) |
| else: |
| code_list = cached_codes[inst_id_to_label_id(i)] |
| random_code = random.choice(code_list) |
| mean_feat = random_code.expand_as(codes[indices[:,0] + b, :, indices[:,2], indices[:,3]]) |
| valid_ratio = torch.ones((1,)) |
| mean_feat = mean_feat * torch.sigmoid(self.gamma(valid_ratio)) + torch.sigmoid(self.beta(valid_ratio)) |
| style_codes[indices[:,0] + b, :, indices[:,2], indices[:,3]] = mean_feat |
| return style_codes |
|
|
| def _forward_layers(self, input): |
| |
| e1 = self.enc1(input) |
| e2 = self.enc2(e1) |
| e3 = self.enc3(e2) |
| e4 = self.enc4(e3) |
| e5 = self.enc5(e4) |
| x = self.enc6(e5) |
| |
| x = F.interpolate(x, scale_factor=2, mode='bilinear') |
| x = torch.cat((x, e4), dim=1) |
| x = self.dec5(x) |
| |
| x = F.interpolate(x, scale_factor=2, mode='bilinear') |
| x = torch.cat((x, e3), dim=1) |
| x = self.dec4(x) |
| |
| x = F.interpolate(x, scale_factor=2, mode='bilinear') |
| x = torch.cat((x, e2), dim=1) |
| x = self.dec3(x) |
| |
| x = F.interpolate(x, scale_factor=2, mode='bilinear') |
| x = torch.cat((x, e1), dim=1) |
| x = self.dec2(x) |
| |
| x = self.dec1(x) |
|
|
| return x |
|
|
|
|
