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