model/shader.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from .warplayer import warp_features
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class DecoderBlock(nn.Module):
    def __init__(self, in_planes, c=224, out_msgs=0, out_locals=0, block_nums=1, out_masks=1, out_local_flows=32, out_msgs_flows=32, out_feat_flows=0):

        super(DecoderBlock, self).__init__()
        self.conv0 = nn.Sequential(
            nn.Conv2d(in_planes, c, 3, 2, 1),
            nn.PReLU(c),
            nn.Conv2d(c, c, 3, 2, 1),
            nn.PReLU(c),
        )

        self.convblocks = nn.ModuleList()
        for i in range(block_nums):
            self.convblocks.append(nn.Sequential(
                nn.Conv2d(c, c, 3, 1, 1),
                nn.PReLU(c),
                nn.Conv2d(c, c, 3, 1, 1),
                nn.PReLU(c),
                nn.Conv2d(c, c, 3, 1, 1),
                nn.PReLU(c),
                nn.Conv2d(c, c, 3, 1, 1),
                nn.PReLU(c),
                nn.Conv2d(c, c, 3, 1, 1),
                nn.PReLU(c),
                nn.Conv2d(c, c, 3, 1, 1),
                nn.PReLU(c),
            ))
        self.out_flows = 2
        self.out_msgs = out_msgs
        self.out_msgs_flows = out_msgs_flows if out_msgs > 0 else 0
        self.out_locals = out_locals
        self.out_local_flows = out_local_flows if out_locals > 0 else 0
        self.out_masks = out_masks
        self.out_feat_flows = out_feat_flows

        self.conv_last = nn.Sequential(
            nn.ConvTranspose2d(c, c, 4, 2, 1),
            nn.PReLU(c),
            nn.ConvTranspose2d(c, self.out_flows+self.out_msgs+self.out_msgs_flows +
                               self.out_locals+self.out_local_flows+self.out_masks+self.out_feat_flows, 4, 2, 1),
        )

    def forward(self, accumulated_flow, *other):
        x = [accumulated_flow]
        for each in other:
            if each is not None:
                assert(accumulated_flow.shape[-1] == each.shape[-1]), "decoder want {}, but get {}".format(
                    accumulated_flow.shape, each.shape)
                x.append(each)
        feat = self.conv0(torch.cat(x, dim=1))
        for convblock1 in self.convblocks:
            feat = convblock1(feat) + feat
        feat = self.conv_last(feat)
        prev = 0
        flow = feat[:, prev:prev+self.out_flows, :, :]
        prev += self.out_flows
        message = feat[:, prev:prev+self.out_msgs,
                       :, :] if self.out_msgs > 0 else None
        prev += self.out_msgs
        message_flow = feat[:, prev:prev + self.out_msgs_flows,
                            :, :] if self.out_msgs_flows > 0 else None
        prev += self.out_msgs_flows
        local_message = feat[:, prev:prev + self.out_locals,
                             :, :] if self.out_locals > 0 else None
        prev += self.out_locals
        local_message_flow = feat[:, prev:prev+self.out_local_flows,
                                  :, :] if self.out_local_flows > 0 else None
        prev += self.out_local_flows
        mask = torch.sigmoid(
            feat[:, prev:prev+self.out_masks, :, :]) if self.out_masks > 0 else None
        prev += self.out_masks
        feat_flow = feat[:, prev:prev+self.out_feat_flows,
                         :, :] if self.out_feat_flows > 0 else None
        prev += self.out_feat_flows
        return flow, mask, message, message_flow, local_message, local_message_flow, feat_flow


class CINN(nn.Module):
    def __init__(self, DIM_SHADER_REFERENCE, target_feature_chns=[512, 256, 128, 64, 64], feature_chns=[2048, 1024, 512, 256, 64], out_msgs_chn=[2048, 1024, 512, 256, 64, 64], out_locals_chn=[2048, 1024, 512, 256, 64, 0], block_num=[1, 1, 1, 1, 1, 2], block_chn_num=[224, 224, 224, 224, 224, 224]):
        super(CINN, self).__init__()
       
        self.in_msgs_chn = [0, *out_msgs_chn[:-1]]
        self.in_locals_chn = [0, *out_locals_chn[:-1]]

        self.decoder_blocks = nn.ModuleList()
        self.feed_weighted = True
        if self.feed_weighted:
            in_planes = 2+2+DIM_SHADER_REFERENCE*2
        else:
            in_planes = 2+DIM_SHADER_REFERENCE
        for each_target_feature_chns, each_feature_chns, each_out_msgs_chn, each_out_locals_chn, each_in_msgs_chn, each_in_locals_chn, each_block_num, each_block_chn_num in zip(target_feature_chns, feature_chns, out_msgs_chn, out_locals_chn, self.in_msgs_chn, self.in_locals_chn, block_num, block_chn_num):
            self.decoder_blocks.append(
                DecoderBlock(in_planes+each_target_feature_chns+each_feature_chns+each_in_locals_chn+each_in_msgs_chn, c=each_block_chn_num, block_nums=each_block_num, out_msgs=each_out_msgs_chn, out_locals=each_out_locals_chn, out_masks=2+each_out_locals_chn))
        for i in range(len(feature_chns), len(out_locals_chn)):
            #print("append extra block", i, "msg",
            #      out_msgs_chn[i], "local", out_locals_chn[i], "block", block_num[i])
            self.decoder_blocks.append(
                DecoderBlock(in_planes+self.in_msgs_chn[i]+self.in_locals_chn[i], c=block_chn_num[i], block_nums=block_num[i], out_msgs=out_msgs_chn[i], out_locals=out_locals_chn[i], out_masks=2+out_msgs_chn[i], out_feat_flows=0))

    def apply_flow(self, mask, message, message_flow, local_message, local_message_flow, x_reference, accumulated_flow, each_x_reference_features=None, each_x_reference_features_flow=None):
        if each_x_reference_features is not None:
            size_from = each_x_reference_features
        else:
            size_from = x_reference
        f_size = (size_from.shape[2], size_from.shape[3])
        accumulated_flow = self.flow_rescale(
            accumulated_flow, size_from)
        # mask = warp_features(F.interpolate(
        #    mask, size=f_size, mode="bilinear"), accumulated_flow) if mask is not None else None
        mask = F.interpolate(
            mask, size=f_size, mode="bilinear") if mask is not None else None
        message = F.interpolate(
            message, size=f_size, mode="bilinear") if message is not None else None
        message_flow = self.flow_rescale(
            message_flow, size_from) if message_flow is not None else None
        message = warp_features(
            message, message_flow) if message_flow is not None else message

        local_message = F.interpolate(
            local_message, size=f_size, mode="bilinear") if local_message is not None else None
        local_message_flow = self.flow_rescale(
            local_message_flow, size_from) if local_message_flow is not None else None
        local_message = warp_features(
            local_message, local_message_flow) if local_message_flow is not None else local_message

        warp_x_reference = warp_features(F.interpolate(
            x_reference, size=f_size, mode="bilinear"), accumulated_flow)

        each_x_reference_features_flow = self.flow_rescale(
            each_x_reference_features_flow, size_from) if (each_x_reference_features is not None and each_x_reference_features_flow is not None) else None
        warp_each_x_reference_features = warp_features(
            each_x_reference_features, each_x_reference_features_flow) if each_x_reference_features_flow is not None else each_x_reference_features

        return mask, message, local_message, warp_x_reference, accumulated_flow, warp_each_x_reference_features, each_x_reference_features_flow

    def forward(self, x_target_features=[], x_reference=None, x_reference_features=[]):
        y_flow = []
        y_feat_flow = []

        y_local_message = []
        y_warp_x_reference = []
        y_warp_x_reference_features = []

        y_weighted_flow = []
        y_weighted_mask = []
        y_weighted_message = []
        y_weighted_x_reference = []
        y_weighted_x_reference_features = []

        for pyrlevel, ifblock in enumerate(self.decoder_blocks):
            stacked_wref = []
            stacked_feat = []
            stacked_anci = []
            stacked_flow = []
            stacked_mask = []
            stacked_mesg = []
            stacked_locm = []
            stacked_feat_flow = []
            for view_id in range(x_reference.shape[1]):  # NMCHW

                if pyrlevel == 0:
                    # create from zero flow
                    feat_ev = x_reference_features[pyrlevel][:,
                                                             view_id, :, :, :] if pyrlevel < len(x_reference_features) else None

                    accumulated_flow = torch.zeros_like(
                        feat_ev[:, :2, :, :]).to(device)
                    accumulated_feat_flow = torch.zeros_like(
                        feat_ev[:, :32, :, :]).to(device)
                    # domestic inputs
                    warp_x_reference = F.interpolate(x_reference[:, view_id, :, :, :], size=(
                        feat_ev.shape[-2], feat_ev.shape[-1]), mode="bilinear")
                    warp_x_reference_features = feat_ev

                    local_message = None
                    # federated inputs
                    weighted_flow = accumulated_flow if self.feed_weighted else None
                    weighted_wref = warp_x_reference if self.feed_weighted else None
                    weighted_message = None
                else:
                    # resume from last layer
                    accumulated_flow = y_flow[-1][:, view_id, :, :, :]
                    accumulated_feat_flow = y_feat_flow[-1][:,
                                                            view_id, :, :, :] if y_feat_flow[-1] is not None else None
                    # domestic inputs
                    warp_x_reference = y_warp_x_reference[-1][:,
                                                              view_id, :, :, :]
                    warp_x_reference_features = y_warp_x_reference_features[-1][:,
                                                                                view_id, :, :, :] if y_warp_x_reference_features[-1] is not None else None
                    local_message = y_local_message[-1][:, view_id, :,
                                                        :, :] if len(y_local_message) > 0 else None

                    # federated inputs
                    weighted_flow = y_weighted_flow[-1] if self.feed_weighted else None
                    weighted_wref = y_weighted_x_reference[-1] if self.feed_weighted else None
                    weighted_message = y_weighted_message[-1] if len(
                        y_weighted_message) > 0 else None
                scaled_x_target = x_target_features[pyrlevel][:, :, :, :].detach() if pyrlevel < len(
                    x_target_features) else None
                # compute flow
                residual_flow, mask, message, message_flow, local_message, local_message_flow, residual_feat_flow = ifblock(
                    accumulated_flow, scaled_x_target, warp_x_reference, warp_x_reference_features, weighted_flow, weighted_wref, weighted_message, local_message)
                accumulated_flow = residual_flow + accumulated_flow
                accumulated_feat_flow = accumulated_flow

                feat_ev = x_reference_features[pyrlevel+1][:,
                                                           view_id, :, :, :] if pyrlevel+1 < len(x_reference_features) else None
                mask, message, local_message, warp_x_reference, accumulated_flow,  warp_x_reference_features, accumulated_feat_flow = self.apply_flow(
                    mask, message, message_flow, local_message, local_message_flow, x_reference[:, view_id, :, :, :], accumulated_flow, feat_ev, accumulated_feat_flow)
                stacked_flow.append(accumulated_flow)
                if accumulated_feat_flow is not None:
                    stacked_feat_flow.append(accumulated_feat_flow)
                stacked_mask.append(mask)
                if message is not None:
                    stacked_mesg.append(message)
                if local_message is not None:
                    stacked_locm.append(local_message)
                stacked_wref.append(warp_x_reference)
                if warp_x_reference_features is not None:
                    stacked_feat.append(warp_x_reference_features)

            stacked_flow = torch.stack(stacked_flow, dim=1)  # M*NCHW -> NMCHW
            stacked_feat_flow = torch.stack(stacked_feat_flow, dim=1) if len(
                stacked_feat_flow) > 0 else None
            stacked_mask = torch.stack(
                stacked_mask, dim=1)
            
            stacked_mesg = torch.stack(stacked_mesg, dim=1) if len(
                stacked_mesg) > 0 else None
            stacked_locm = torch.stack(stacked_locm, dim=1) if len(
                stacked_locm) > 0 else None

            stacked_wref = torch.stack(stacked_wref, dim=1)
            stacked_feat = torch.stack(stacked_feat, dim=1) if len(
                stacked_feat) > 0 else None
            stacked_anci = torch.stack(stacked_anci, dim=1) if len(
                stacked_anci) > 0 else None
            y_flow.append(stacked_flow)
            y_feat_flow.append(stacked_feat_flow)

            y_warp_x_reference.append(stacked_wref)
            y_warp_x_reference_features.append(stacked_feat)
            # compute normalized confidence
            stacked_contrib = torch.nn.functional.softmax(stacked_mask, dim=1)

            # torch.sum to remove temp dimension M from NMCHW --> NCHW
            weighted_flow = torch.sum(
                stacked_mask[:, :, 0:1, :, :] * stacked_contrib[:, :, 0:1, :, :] * stacked_flow, dim=1)
            weighted_mask = torch.sum(
                stacked_contrib[:, :, 0:1, :, :] * stacked_mask[:, :, 0:1, :, :], dim=1)
            weighted_wref = torch.sum(
                stacked_mask[:, :, 0:1, :, :] * stacked_contrib[:, :, 0:1, :, :] * stacked_wref, dim=1) if stacked_wref is not None else None
            weighted_feat = torch.sum(
                stacked_mask[:, :, 1:2, :, :] * stacked_contrib[:, :, 1:2, :, :] * stacked_feat, dim=1) if stacked_feat is not None else None
            weighted_mesg = torch.sum(
                stacked_mask[:, :, 2:, :, :] * stacked_contrib[:, :, 2:, :, :] * stacked_mesg, dim=1) if stacked_mesg is not None else None
            y_weighted_flow.append(weighted_flow)
            y_weighted_mask.append(weighted_mask)
            if weighted_mesg is not None:
                y_weighted_message.append(weighted_mesg)
            if stacked_locm is not None:
                y_local_message.append(stacked_locm)
            y_weighted_message.append(weighted_mesg)
            y_weighted_x_reference.append(weighted_wref)
            y_weighted_x_reference_features.append(weighted_feat)

            if weighted_feat is not None:
                y_weighted_x_reference_features.append(weighted_feat)
        return {
            "y_last_remote_features": [weighted_mesg],
        }

    def flow_rescale(self, prev_flow, each_x_reference_features):
        if prev_flow is None:
            prev_flow = torch.zeros_like(
                each_x_reference_features[:, :2]).to(device)
        else:
            up_scale_factor = each_x_reference_features.shape[-1] / \
                prev_flow.shape[-1]
            if up_scale_factor != 1:
                prev_flow = F.interpolate(prev_flow, scale_factor=up_scale_factor, mode="bilinear",
                                          align_corners=False, recompute_scale_factor=False) * up_scale_factor
        return prev_flow