models/rs_super.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat

import models
from models import register
from utils import make_coord, to_coordinates


@register('rs_super')
class RSSuper(nn.Module):
    def __init__(self,
                 encoder_spec,
                 neck=None,
                 decoder=None,
                 input_rgb=True,
                 n_forward_times=1,
                 global_decoder=None
                 ):
        super().__init__()
        self.n_forward_times = n_forward_times
        self.encoder = models.make(encoder_spec)
        if neck is not None:
            self.neck = models.make(neck, args={'in_dim': self.encoder.out_dim})

        self.input_rgb = input_rgb
        decoder_in_dim = 5 if self.input_rgb else 2
        if decoder is not None:
            self.decoder = models.make(decoder, args={'modulation_dim': self.neck.out_dim, 'in_dim': decoder_in_dim})

        if global_decoder is not None:
            decoder_in_dim = 5 if self.input_rgb else 2
            self.decoder_is_proj = global_decoder.get('is_proj', False)
            self.grid_global = global_decoder.get('grid_global', False)

            self.global_decoder = models.make(global_decoder, args={'modulation_dim': self.neck.out_dim, 'in_dim': decoder_in_dim})

            if self.decoder_is_proj:
                self.input_proj = nn.Sequential(
                    nn.Linear(self.neck.out_dim, self.neck.out_dim)
                    )
                self.output_proj = nn.Sequential(
                    nn.Linear(3, 3)
                    )

    def query_rgb(self, coord, cell=None):
        feat = self.feat

        if self.imnet is None:
            ret = F.grid_sample(feat, coord.flip(-1).unsqueeze(1),
                                mode='nearest', align_corners=False)[:, :, 0, :] \
                .permute(0, 2, 1)
            return ret

        if self.feat_unfold:
            feat = F.unfold(feat, 3, padding=1).view(
                feat.shape[0], feat.shape[1] * 9, feat.shape[2], feat.shape[3])

        if self.local_ensemble:
            vx_lst = [-1, 1]
            vy_lst = [-1, 1]
            eps_shift = 1e-6
        else:
            vx_lst, vy_lst, eps_shift = [0], [0], 0

        # field radius (global: [-1, 1])
        rx = 2 / feat.shape[-2] / 2
        ry = 2 / feat.shape[-1] / 2

        feat_coord = make_coord(feat.shape[-2:], flatten=False).cuda() \
            .permute(2, 0, 1) \
            .unsqueeze(0).expand(feat.shape[0], 2, *feat.shape[-2:])

        preds = []
        areas = []
        for vx in vx_lst:
            for vy in vy_lst:
                coord_ = coord.clone()
                coord_[:, :, 0] += vx * rx + eps_shift
                coord_[:, :, 1] += vy * ry + eps_shift
                coord_.clamp_(-1 + 1e-6, 1 - 1e-6)
                q_feat = F.grid_sample(
                    feat, coord_.flip(-1).unsqueeze(1),
                    mode='nearest', align_corners=False)[:, :, 0, :] \
                    .permute(0, 2, 1)
                q_coord = F.grid_sample(
                    feat_coord, coord_.flip(-1).unsqueeze(1),
                    mode='nearest', align_corners=False)[:, :, 0, :] \
                    .permute(0, 2, 1)
                rel_coord = coord - q_coord
                rel_coord[:, :, 0] *= feat.shape[-2]
                rel_coord[:, :, 1] *= feat.shape[-1]
                inp = torch.cat([q_feat, rel_coord], dim=-1)

                if self.cell_decode:
                    rel_cell = cell.clone()
                    rel_cell[:, :, 0] *= feat.shape[-2]
                    rel_cell[:, :, 1] *= feat.shape[-1]
                    inp = torch.cat([inp, rel_cell], dim=-1)

                bs, q = coord.shape[:2]
                pred = self.imnet(inp.view(bs * q, -1)).view(bs, q, -1)
                preds.append(pred)

                area = torch.abs(rel_coord[:, :, 0] * rel_coord[:, :, 1])
                areas.append(area + 1e-9)

        tot_area = torch.stack(areas).sum(dim=0)
        if self.local_ensemble:
            t = areas[0]; areas[0] = areas[3]; areas[3] = t
            t = areas[1]; areas[1] = areas[2]; areas[2] = t
        ret = 0
        for pred, area in zip(preds, areas):
            ret = ret + pred * (area / tot_area).unsqueeze(-1)
        return ret

    def forward_backbone_neck(self, inp, coord):
        # inp: 64x3x32x32
        # coord: BxNx2
        feat = self.encoder(inp)  # 64x64x32x32
        global_content, x_rep = self.neck(feat)  # Bx1xC; BxCxHxW
        return feat, x_rep, global_content

    def forward_step(self, inp, coord, feat, x_rep, global_content, pred_rgb_value=None):
        weight_gen_func = 'bilinear'  # 'bilinear'
        #  grid： 先x再y
        coord_ = coord.clone().unsqueeze(1).flip(-1)  # Bx1xNxC
        modulations = F.grid_sample(x_rep, coord_, padding_mode='border', mode=weight_gen_func,
                                    align_corners=True).squeeze(2)  # B C N
        modulations = rearrange(modulations, 'B C N -> (B N) C')

        feat_coord = to_coordinates(feat.shape[-2:], return_map=True).to(inp.device)
        feat_coord = repeat(feat_coord, 'H W C -> B C H W', B=inp.size(0))  # 坐标是[y, x]
        nearest_coord = F.grid_sample(feat_coord, coord_, mode='nearest', align_corners=True).squeeze(2)  # B 2 N
        nearest_coord = rearrange(nearest_coord, 'B C N -> B N C')  # B N 2

        relative_coord = coord - nearest_coord
        relative_coord[:, :, 0] *= feat.shape[-2]
        relative_coord[:, :, 1] *= feat.shape[-1]
        relative_coord = rearrange(relative_coord, 'B N C -> (B N) C')
        decoder_input = relative_coord

        interpolated_rgb = None
        if self.input_rgb:
            if pred_rgb_value is not None:
                interpolated_rgb = rearrange(pred_rgb_value, 'B N C -> (B N) C')
            else:
                interpolated_rgb = F.grid_sample(inp, coord_, padding_mode='border', mode='bilinear', align_corners=True).squeeze(2)  # B 3 N
                interpolated_rgb = rearrange(interpolated_rgb, 'B C N -> (B N) C')
            decoder_input = torch.cat((decoder_input, interpolated_rgb), dim=-1)

        decoder_output = self.decoder(decoder_input, modulations)
        decoder_output = rearrange(decoder_output, '(B N) C -> B N C', B=inp.size(0))

        if hasattr(self, 'global_decoder'):
            # coord: BxNx2
            # global_content: Bx1xC
            if self.decoder_is_proj:
                global_content = self.input_proj(global_content)  # B 1 C
            global_modulations = repeat(global_content, 'B N C -> B (N S) C', S=coord.size(1))
            global_modulations = rearrange(global_modulations, 'B N C -> (B N) C')

            if self.grid_global:
                # import pdb
                # pdb.set_trace()
                global_decoder_input = decoder_input
            else:
                global_decoder_input = rearrange(coord, 'B N C -> (B N) C')
                if self.input_rgb:
                    global_decoder_input = torch.cat((global_decoder_input, interpolated_rgb), dim=-1)

            global_decoder_output = self.global_decoder(global_decoder_input, global_modulations)
            global_decoder_output = rearrange(global_decoder_output, '(B N) C -> B N C', B=inp.size(0))

            if self.decoder_is_proj:
                decoder_output = self.output_proj(global_decoder_output + decoder_output)
            else:
                decoder_output = global_decoder_output + decoder_output

        return decoder_output

    def forward(self, inp, coord):
        # import pdb
        # pdb.set_trace()
        pred_rgb_value = None
        feat, x_rep, global_content = self.forward_backbone_neck(inp, coord)
        return_pred_rgb_value = []
        for n_time in range(self.n_forward_times):
            pred_rgb_value = self.forward_step(inp, coord, feat, x_rep, global_content, pred_rgb_value)
            return_pred_rgb_value.append(pred_rgb_value)
        return return_pred_rgb_value