Spaces:
Sleeping
Sleeping
| import torch | |
| from torch import nn | |
| import torch.nn.functional as F | |
| from modules.util import ResBlock2d, SameBlock2d, UpBlock2d, DownBlock2d | |
| from modules.dense_motion import DenseMotionNetwork | |
| from modules.nerf_verts_util import RenderModel | |
| class SPADE_layer(nn.Module): | |
| def __init__(self, norm_channel, label_channel): | |
| super(SPADE_layer, self).__init__() | |
| self.param_free_norm = nn.InstanceNorm2d(norm_channel, affine=False) | |
| hidden_channel = 128 | |
| self.mlp_shared = nn.Sequential( | |
| nn.Conv2d(label_channel, hidden_channel, kernel_size=3, padding=1), | |
| nn.ReLU() | |
| ) | |
| self.mlp_gamma = nn.Conv2d(hidden_channel, norm_channel, kernel_size=3, padding=1) | |
| self.mlp_beta = nn.Conv2d(hidden_channel, norm_channel, kernel_size=3, padding=1) | |
| def forward(self, x, modulation_in): | |
| normalized = self.param_free_norm(x) | |
| modulation_in = F.interpolate(modulation_in, size=x.size()[2:], mode='nearest') | |
| actv = self.mlp_shared(modulation_in) | |
| gamma = self.mlp_gamma(actv) | |
| beta = self.mlp_beta(actv) | |
| out = normalized * (1 + gamma) + beta | |
| return out | |
| class SPADE_block(nn.Module): | |
| def __init__(self, norm_channel, label_channel, out_channel): | |
| super(SPADE_block, self).__init__() | |
| self.SPADE_0 = SPADE_layer(norm_channel, label_channel) | |
| self.relu_0 = nn.ReLU() | |
| self.conv_0 = nn.Conv2d(norm_channel, norm_channel, kernel_size=3, padding=1) | |
| self.SPADE_1 = SPADE_layer(norm_channel, label_channel) | |
| self.relu_1 = nn.ReLU() | |
| self.conv_1 = nn.Conv2d(norm_channel, out_channel, kernel_size=3, padding=1) | |
| def forward(self, x, modulation_in): | |
| out = self.SPADE_0(x, modulation_in) | |
| out = self.relu_0(out) | |
| out = self.conv_0(out) | |
| out = self.SPADE_1(out, modulation_in) | |
| out = self.relu_1(out) | |
| out = self.conv_1(out) | |
| return out | |
| class SPADE_decoder(nn.Module): | |
| def __init__(self, in_channel, mid_channel): | |
| super(SPADE_decoder, self).__init__() | |
| self.in_channel = in_channel | |
| self.mid_channel = mid_channel | |
| self.seg_conv = nn.Sequential( | |
| nn.Conv2d(in_channel, mid_channel, kernel_size=3, padding=1), | |
| nn.ReLU() | |
| ) | |
| self.SPADE_0 = SPADE_block(in_channel, mid_channel, in_channel // 4) | |
| self.up_0 = nn.UpsamplingBilinear2d(scale_factor=2) | |
| in_channel = in_channel // 4 | |
| self.SPADE_1 = SPADE_block(in_channel, mid_channel, in_channel // 4) | |
| self.up_1 = nn.UpsamplingBilinear2d(scale_factor=2) | |
| in_channel = in_channel // 4 | |
| self.SPADE_2 = SPADE_block(in_channel, mid_channel, in_channel) | |
| self.SPADE_3 = SPADE_block(in_channel, mid_channel, in_channel) | |
| self.final = nn.Sequential( | |
| nn.Conv2d(in_channel, 3, kernel_size=7, padding=3), | |
| nn.Sigmoid() | |
| ) | |
| def forward(self, x): | |
| seg = self.seg_conv(x) | |
| x = self.SPADE_0(x, seg) | |
| x = self.up_0(x) | |
| x = self.SPADE_1(x, seg) | |
| x = self.up_1(x) | |
| x = self.SPADE_2(x, seg) | |
| x = self.SPADE_3(x, seg) | |
| x = self.final(x) | |
| return x | |
| def calc_mean_std(feat, eps=1e-5): | |
| # eps is a small value added to the variance to avoid divide-by-zero. | |
| size = feat.size() | |
| assert (len(size) == 4) | |
| N, C = size[:2] | |
| feat_var = feat.view(N, C, -1).var(dim=2) + eps | |
| feat_std = feat_var.sqrt().view(N, C, 1, 1) | |
| feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1) | |
| return feat_mean, feat_std | |
| def adaptive_instance_normalization(x, modulation_in): | |
| assert (x.size()[:2] == modulation_in.size()[:2]) | |
| size = x.size() | |
| style_mean, style_std = calc_mean_std(modulation_in) | |
| content_mean, content_std = calc_mean_std(x) | |
| normalized_feat = (x - content_mean.expand( | |
| size)) / content_std.expand(size) | |
| return normalized_feat * style_std.expand(size) + style_mean.expand(size) | |
| class AdaIN_layer(nn.Module): | |
| def __init__(self, norm_channel, label_channel): | |
| super(AdaIN_layer, self).__init__() | |
| self.param_free_norm = nn.InstanceNorm2d(norm_channel, affine=False) | |
| self.mlp_shared = nn.Sequential( | |
| nn.Conv2d(label_channel, norm_channel, kernel_size=3, padding=1), | |
| nn.ReLU() | |
| ) | |
| def forward(self, x, modulation_in): | |
| normalized = self.param_free_norm(x) | |
| modulation_in = self.mlp_shared(modulation_in) | |
| out = adaptive_instance_normalization(normalized, modulation_in) | |
| return out | |
| class OcclusionAwareGenerator_SPADE(nn.Module): | |
| """ | |
| Generator that given source image and and keypoints try to transform image according to movement trajectories | |
| induced by keypoints. Generator follows Johnson architecture. | |
| """ | |
| def __init__(self, num_channels, num_kp, block_expansion, max_features, num_down_blocks, | |
| num_bottleneck_blocks, estimate_occlusion_map=False, dense_motion_params=None, render_params=None, | |
| estimate_jacobian=False): | |
| super(OcclusionAwareGenerator_SPADE, self).__init__() | |
| if dense_motion_params is not None: | |
| self.dense_motion_network = DenseMotionNetwork(num_kp=num_kp, num_channels=num_channels, | |
| estimate_occlusion_map=estimate_occlusion_map, | |
| **dense_motion_params) | |
| else: | |
| self.dense_motion_network = None | |
| self.first = SameBlock2d(num_channels, block_expansion, kernel_size=(7, 7), padding=(3, 3)) | |
| down_blocks = [] | |
| for i in range(num_down_blocks): | |
| in_features = min(max_features, block_expansion * (2 ** i)) | |
| out_features = min(max_features, block_expansion * (2 ** (i + 1))) | |
| down_blocks.append(DownBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1))) | |
| self.down_blocks = nn.ModuleList(down_blocks) | |
| in_features = min(max_features, block_expansion * (2 ** num_down_blocks)) | |
| self.Render_model = RenderModel(in_channels=in_features, **render_params) | |
| self.decoder = SPADE_decoder(in_channel=in_features * 2, mid_channel=128) | |
| self.estimate_occlusion_map = estimate_occlusion_map | |
| self.num_channels = num_channels | |
| def deform_input(self, inp, deformation): | |
| _, h_old, w_old, _ = deformation.shape | |
| _, _, h, w = inp.shape | |
| if h_old != h or w_old != w: | |
| deformation = deformation.permute(0, 3, 1, 2) | |
| deformation = F.interpolate(deformation, size=(h, w), mode='bilinear') | |
| deformation = deformation.permute(0, 2, 3, 1) | |
| return F.grid_sample(inp, deformation) | |
| def forward(self, source_image, kp_driving, kp_source): | |
| # Encoding (downsampling) part | |
| out = self.first(source_image) | |
| for i in range(len(self.down_blocks)): | |
| out = self.down_blocks[i](out) | |
| # Transforming feature representation according to deformation and occlusion | |
| output_dict = {} | |
| if self.dense_motion_network is not None: | |
| dense_motion = self.dense_motion_network(source_image=source_image, kp_driving=kp_driving, | |
| kp_source=kp_source) | |
| output_dict['mask'] = dense_motion['mask'] | |
| output_dict['sparse_deformed'] = dense_motion['sparse_deformed'] | |
| if 'occlusion_map' in dense_motion: | |
| occlusion_map = dense_motion['occlusion_map'] | |
| output_dict['occlusion_map'] = occlusion_map | |
| else: | |
| occlusion_map = None | |
| deformation = dense_motion['deformation'] | |
| out = self.deform_input(out, deformation) | |
| if occlusion_map is not None: | |
| if out.shape[2] != occlusion_map.shape[2] or out.shape[3] != occlusion_map.shape[3]: | |
| occlusion_map = F.interpolate(occlusion_map, size=out.shape[2:], mode='bilinear') | |
| out = out * occlusion_map | |
| output_dict["deformed"] = self.deform_input(source_image, deformation) | |
| # render part | |
| render_result = self.Render_model(feature=out) | |
| output_dict['render'] = render_result['mini_pred'] | |
| output_dict['point_pred'] = render_result['point_pred'] | |
| out = torch.cat((out, render_result['render']), dim=1) | |
| # out = self.merge_conv(out) | |
| # Decoding part | |
| out = self.decoder(out) | |
| output_dict["prediction"] = out | |
| return output_dict | |