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 from mmcv.cnn import ConvModule from .blocks.CSPLayer import CSPLayer @register('rs_multiscale_super') class RSMultiScaleSuper(nn.Module): def __init__(self, encoder_spec, multiscale=False, neck=None, decoder=None, has_bn=True, input_rgb=False, n_forward_times=1, global_decoder=None, encode_scale_ratio=False ): super().__init__() self.encoder = models.make(encoder_spec) self.multiscale = multiscale self.encoder_out_dim = self.encoder.out_dim self.encode_scale_ratio = encode_scale_ratio conv_cfg = None if has_bn: norm_cfg = dict(type='BN', momentum=0.03, eps=0.001) else: norm_cfg = None act_cfg = dict(type='ReLU') if self.multiscale: self.multiscale_layers = nn.ModuleList() # 32->16->8->4 num_blocks = [2, 4, 6] for n_idx in range(3): conv_layer = ConvModule( self.encoder.out_dim, self.encoder.out_dim*2, 3, stride=2, padding=1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg ) csp_layer = CSPLayer( self.encoder.out_dim*2, self.encoder.out_dim, num_blocks=num_blocks[n_idx], add_identity=True, use_depthwise=False, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg) self.multiscale_layers.append(nn.Sequential(conv_layer, csp_layer)) if neck is not None: self.neck = models.make(neck, args={'in_dim': self.encoder.out_dim}) modulation_dim = self.neck.d_dim else: modulation_dim = self.encoder.out_dim self.n_forward_times = n_forward_times self.input_rgb = input_rgb decoder_in_dim = 5 if self.input_rgb else 2 if encode_scale_ratio: decoder_in_dim += 2 if decoder is not None: self.decoder = models.make(decoder, args={'modulation_dim': modulation_dim, 'in_dim': decoder_in_dim}) if global_decoder is not None: decoder_in_dim = 5 if self.input_rgb else 2 if encode_scale_ratio: decoder_in_dim += 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': modulation_dim, 'in_dim': decoder_in_dim}) if self.decoder_is_proj: self.input_proj = nn.Sequential( nn.Linear(modulation_dim, modulation_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: # if kernel_size=2 and stride=2, x should has shape (B, 4*C, H/2*W/2) 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_step(self, ori_img, coord, func_map, global_content, pred_rgb_value=None, scale_ratio=None ): weight_gen_func = 'bilinear' # 'bilinear' # grid： 先x再y coord_ = coord.clone().unsqueeze(1).flip(-1) # Bx1xNxC funcs = F.grid_sample( func_map, coord_, padding_mode='border', mode=weight_gen_func, align_corners=True).squeeze(2) # B C N funcs = rearrange(funcs, 'B C N -> (B N) C') feat_coord = to_coordinates(func_map.shape[-2:], return_map=True).to(func_map.device) feat_coord = repeat(feat_coord, 'H W C -> B C H W', B=coord.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] *= func_map.shape[-2] relative_coord[:, :, 1] *= func_map.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( ori_img, 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) if self.encode_scale_ratio: scale_ratio = rearrange(scale_ratio, 'B N C -> (B N) C') decoder_input = torch.cat((decoder_input, scale_ratio), dim=-1) decoder_output = self.decoder(decoder_input, funcs) decoder_output = rearrange(decoder_output, '(B N) C -> B N C', B=func_map.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_funcs = repeat(global_content, 'B C -> B N C', N=coord.size(1)) global_funcs = rearrange(global_funcs, 'B N C -> (B N) C') if self.grid_global: 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) if self.encode_scale_ratio: global_decoder_input = torch.cat((global_decoder_input, scale_ratio), dim=-1) global_decoder_output = self.global_decoder(global_decoder_input, global_funcs) global_decoder_output = rearrange(global_decoder_output, '(B N) C -> B N C', B=func_map.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_backbone(self, inp): # inp: img-BxCxHxW return self.encoder(inp) def forward_multiscale(self, feats, keep_ori_featmap=False): if keep_ori_featmap: output_feats = feats else: output_feats = [] x = feats[0] for layer in self.multiscale_layers: x = layer(x) output_feats.append(x) return output_feats def forward(self, inp, coord, scale_ratio=None): output_feats = [self.forward_backbone(inp)] if self.multiscale: output_feats = self.forward_multiscale(output_feats) if hasattr(self, 'neck'): global_content, func_maps = self.neck(output_feats) else: global_content = None func_maps = output_feats[0] pred_rgb_value = None return_pred_rgb_value = [] for n_time in range(self.n_forward_times): pred_rgb_value = self.forward_step(inp, coord, func_maps, global_content, pred_rgb_value, scale_ratio) return_pred_rgb_value.append(pred_rgb_value) return return_pred_rgb_value