climategan/masker.py

import torch
import torch.nn as nn
import torch.nn.functional as F

from climategan.blocks import (
    BaseDecoder,
    Conv2dBlock,
    InterpolateNearest2d,
    SPADEResnetBlock,
)


def create_mask_decoder(opts, no_init=False, verbose=0):
    if opts.gen.m.use_spade:
        if verbose > 0:
            print("  - Add Spade Mask Decoder")
        assert "d" in opts.tasks or "s" in opts.tasks
        return MaskSpadeDecoder(opts)
    else:
        if verbose > 0:
            print("  - Add Base Mask Decoder")
        return MaskBaseDecoder(opts)


class MaskBaseDecoder(BaseDecoder):
    def __init__(self, opts):
        low_level_feats_dim = -1
        use_v3 = opts.gen.encoder.architecture == "deeplabv3"
        use_mobile_net = opts.gen.deeplabv3.backbone == "mobilenet"
        use_low = opts.gen.m.use_low_level_feats
        use_dada = ("d" in opts.tasks) and opts.gen.m.use_dada

        if use_v3 and use_mobile_net:
            input_dim = 320
            if use_low:
                low_level_feats_dim = 24
        elif use_v3:
            input_dim = 2048
            if use_low:
                low_level_feats_dim = 256
        else:
            input_dim = 2048

        super().__init__(
            n_upsample=opts.gen.m.n_upsample,
            n_res=opts.gen.m.n_res,
            input_dim=input_dim,
            proj_dim=opts.gen.m.proj_dim,
            output_dim=opts.gen.m.output_dim,
            norm=opts.gen.m.norm,
            activ=opts.gen.m.activ,
            pad_type=opts.gen.m.pad_type,
            output_activ="none",
            low_level_feats_dim=low_level_feats_dim,
            use_dada=use_dada,
        )


class MaskSpadeDecoder(nn.Module):
    def __init__(self, opts):
        """Create a SPADE-based decoder, which forwards z and the conditioning
        tensors seg (in the original paper, conditioning is on a semantic map only).
        All along, z is conditioned on seg. First 3 SpadeResblocks (SRB) do not shrink
        the channel dimension, and an upsampling is applied after each. Therefore
        2 upsamplings at this point. Then, for each remaining upsamplings
        (w.r.t. spade_n_up), the SRB shrinks channels by 2. Before final conv to get 3
        channels, the number of channels is therefore:
            final_nc = channels(z) * 2 ** (spade_n_up - 2)
        Args:
            latent_dim (tuple): z's shape (only the number of channels matters)
            cond_nc (int): conditioning tensor's expected number of channels
            spade_n_up (int): Number of total upsamplings from z
            spade_use_spectral_norm (bool): use spectral normalization?
            spade_param_free_norm (str): norm to use before SPADE de-normalization
            spade_kernel_size (int): SPADE conv layers' kernel size
        Returns:
            [type]: [description]
        """
        super().__init__()
        self.opts = opts
        latent_dim = opts.gen.m.spade.latent_dim
        cond_nc = opts.gen.m.spade.cond_nc
        spade_use_spectral_norm = opts.gen.m.spade.spade_use_spectral_norm
        spade_param_free_norm = opts.gen.m.spade.spade_param_free_norm
        if self.opts.gen.m.spade.activations.all_lrelu:
            spade_activation = "lrelu"
        else:
            spade_activation = None
        spade_kernel_size = 3
        self.num_layers = opts.gen.m.spade.num_layers
        self.z_nc = latent_dim

        if (
            opts.gen.encoder.architecture == "deeplabv3"
            and opts.gen.deeplabv3.backbone == "mobilenet"
        ):
            self.input_dim = [320, 24]
            self.low_level_conv = Conv2dBlock(
                self.input_dim[1],
                self.input_dim[0],
                3,
                padding=1,
                activation="lrelu",
                pad_type="reflect",
                norm="spectral_batch",
            )
            self.merge_feats_conv = Conv2dBlock(
                self.input_dim[0] * 2,
                self.z_nc,
                3,
                padding=1,
                activation="lrelu",
                pad_type="reflect",
                norm="spectral_batch",
            )
        elif (
            opts.gen.encoder.architecture == "deeplabv3"
            and opts.gen.deeplabv3.backbone == "resnet"
        ):
            self.input_dim = [2048, 256]
            if self.opts.gen.m.use_proj:
                proj_dim = self.opts.gen.m.proj_dim
                self.low_level_conv = Conv2dBlock(
                    self.input_dim[1],
                    proj_dim,
                    3,
                    padding=1,
                    activation="lrelu",
                    pad_type="reflect",
                    norm="spectral_batch",
                )
                self.high_level_conv = Conv2dBlock(
                    self.input_dim[0],
                    proj_dim,
                    3,
                    padding=1,
                    activation="lrelu",
                    pad_type="reflect",
                    norm="spectral_batch",
                )
                self.merge_feats_conv = Conv2dBlock(
                    proj_dim * 2,
                    self.z_nc,
                    3,
                    padding=1,
                    activation="lrelu",
                    pad_type="reflect",
                    norm="spectral_batch",
                )
            else:
                self.low_level_conv = Conv2dBlock(
                    self.input_dim[1],
                    self.input_dim[0],
                    3,
                    padding=1,
                    activation="lrelu",
                    pad_type="reflect",
                    norm="spectral_batch",
                )
                self.merge_feats_conv = Conv2dBlock(
                    self.input_dim[0] * 2,
                    self.z_nc,
                    3,
                    padding=1,
                    activation="lrelu",
                    pad_type="reflect",
                    norm="spectral_batch",
                )

        elif opts.gen.encoder.architecture == "deeplabv2":
            self.input_dim = 2048
            self.fc_conv = Conv2dBlock(
                self.input_dim,
                self.z_nc,
                3,
                padding=1,
                activation="lrelu",
                pad_type="reflect",
                norm="spectral_batch",
            )
        else:
            raise ValueError("Unknown encoder type")

        self.spade_blocks = []

        for i in range(self.num_layers):
            self.spade_blocks.append(
                SPADEResnetBlock(
                    int(self.z_nc / (2**i)),
                    int(self.z_nc / (2 ** (i + 1))),
                    cond_nc,
                    spade_use_spectral_norm,
                    spade_param_free_norm,
                    spade_kernel_size,
                    spade_activation,
                )
            )
        self.spade_blocks = nn.Sequential(*self.spade_blocks)

        self.final_nc = int(self.z_nc / (2**self.num_layers))
        self.mask_conv = Conv2dBlock(
            self.final_nc,
            1,
            3,
            padding=1,
            activation="none",
            pad_type="reflect",
            norm="spectral",
        )
        self.upsample = InterpolateNearest2d(scale_factor=2)

    def forward(self, z, cond, z_depth=None):
        if isinstance(z, (list, tuple)):
            z_h, z_l = z
            if self.opts.gen.m.use_proj:
                z_l = self.low_level_conv(z_l)
                z_l = F.interpolate(z_l, size=z_h.shape[-2:], mode="bilinear")
                z_h = self.high_level_conv(z_h)
            else:
                z_l = self.low_level_conv(z_l)
                z_l = F.interpolate(z_l, size=z_h.shape[-2:], mode="bilinear")
            z = torch.cat([z_h, z_l], axis=1)
            y = self.merge_feats_conv(z)
        else:
            y = self.fc_conv(z)

        for i in range(self.num_layers):
            y = self.spade_blocks[i](y, cond)
            y = self.upsample(y)
        y = self.mask_conv(y)
        return y

    def __str__(self):
        return "MaskerSpadeDecoder"