modules/generator.py

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