import torch import torch.nn as nn import torch.nn.functional as F from .swin_transformer import SwinTransformer from .newcrf_layers import NewCRF from .uper_crf_head import PSP from .depth_update import * ######################################################################################################################## class NewCRFDepth(nn.Module): """ Depth network based on neural window FC-CRFs architecture. """ def __init__(self, version=None, inv_depth=False, pretrained=None, frozen_stages=-1, min_depth=0.1, max_depth=100.0, **kwargs): super().__init__() self.inv_depth = inv_depth self.with_auxiliary_head = False self.with_neck = False norm_cfg = dict(type='BN', requires_grad=True) window_size = int(version[-2:]) if version[:-2] == 'base': embed_dim = 128 depths = [2, 2, 18, 2] num_heads = [4, 8, 16, 32] in_channels = [128, 256, 512, 1024] self.update = BasicUpdateBlockDepth(hidden_dim=128, context_dim=128) elif version[:-2] == 'large': embed_dim = 192 depths = [2, 2, 18, 2] num_heads = [6, 12, 24, 48] in_channels = [192, 384, 768, 1536] self.update = BasicUpdateBlockDepth(hidden_dim=128, context_dim=192) elif version[:-2] == 'tiny': embed_dim = 96 depths = [2, 2, 6, 2] num_heads = [3, 6, 12, 24] in_channels = [96, 192, 384, 768] self.update = BasicUpdateBlockDepth(hidden_dim=128, context_dim=96) backbone_cfg = dict( embed_dim=embed_dim, depths=depths, num_heads=num_heads, window_size=window_size, ape=False, drop_path_rate=0.3, patch_norm=True, use_checkpoint=False, frozen_stages=frozen_stages ) embed_dim = 512 decoder_cfg = dict( in_channels=in_channels, in_index=[0, 1, 2, 3], pool_scales=(1, 2, 3, 6), channels=embed_dim, dropout_ratio=0.0, num_classes=32, norm_cfg=norm_cfg, align_corners=False ) self.backbone = SwinTransformer(**backbone_cfg) v_dim = decoder_cfg['num_classes']*4 win = 7 crf_dims = [128, 256, 512, 1024] v_dims = [64, 128, 256, embed_dim] self.crf3 = NewCRF(input_dim=in_channels[3], embed_dim=crf_dims[3], window_size=win, v_dim=v_dims[3], num_heads=32) self.crf2 = NewCRF(input_dim=in_channels[2], embed_dim=crf_dims[2], window_size=win, v_dim=v_dims[2], num_heads=16) self.crf1 = NewCRF(input_dim=in_channels[1], embed_dim=crf_dims[1], window_size=win, v_dim=v_dims[1], num_heads=8) self.decoder = PSP(**decoder_cfg) self.disp_head1 = DispHead(input_dim=crf_dims[0]) self.up_mode = 'bilinear' if self.up_mode == 'mask': self.mask_head = nn.Sequential( nn.Conv2d(v_dims[0], 64, 3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(64, 16*9, 1, padding=0)) self.min_depth = min_depth self.max_depth = max_depth self.depth_num = 16 self.hidden_dim = 128 self.project = Projection(v_dims[0], self.hidden_dim) self.init_weights(pretrained=pretrained) def init_weights(self, pretrained=None): """Initialize the weights in backbone and heads. Args: pretrained (str, optional): Path to pre-trained weights. Defaults to None. """ print(f'== Load encoder backbone from: {pretrained}') self.backbone.init_weights(pretrained=pretrained) self.decoder.init_weights() if self.with_auxiliary_head: if isinstance(self.auxiliary_head, nn.ModuleList): for aux_head in self.auxiliary_head: aux_head.init_weights() else: self.auxiliary_head.init_weights() def upsample_mask(self, disp, mask): """ Upsample disp [H/4, W/4, 1] -> [H, W, 1] using convex combination """ N, C, H, W = disp.shape mask = mask.view(N, 1, 9, 4, 4, H, W) mask = torch.softmax(mask, dim=2) up_disp = F.unfold(disp, kernel_size=3, padding=1) up_disp = up_disp.view(N, C, 9, 1, 1, H, W) up_disp = torch.sum(mask * up_disp, dim=2) up_disp = up_disp.permute(0, 1, 4, 2, 5, 3) return up_disp.reshape(N, C, 4*H, 4*W) def forward(self, imgs, epoch=1, step=100): feats = self.backbone(imgs) ppm_out = self.decoder(feats) e3 = self.crf3(feats[3], ppm_out) e3 = nn.PixelShuffle(2)(e3) e2 = self.crf2(feats[2], e3) e2 = nn.PixelShuffle(2)(e2) e1 = self.crf1(feats[1], e2) e1 = nn.PixelShuffle(2)(e1) # iterative bins if epoch == 0 and step < 80: max_tree_depth = 3 else: max_tree_depth = 6 if self.up_mode == 'mask': mask = self.mask_head(e1) b, c, h, w = e1.shape device = e1.device depth = torch.zeros([b, 1, h, w]).to(device) context = feats[0] gru_hidden = torch.tanh(self.project(e1)) pred_depths_r_list, pred_depths_c_list, uncertainty_maps_list = self.update(depth, context, gru_hidden, max_tree_depth, self.depth_num, self.min_depth, self.max_depth) if self.up_mode == 'mask': for i in range(len(pred_depths_r_list)): pred_depths_r_list[i] = self.upsample_mask(pred_depths_r_list[i], mask) for i in range(len(pred_depths_c_list)): pred_depths_c_list[i] = self.upsample_mask(pred_depths_c_list[i], mask.detach()) for i in range(len(uncertainty_maps_list)): uncertainty_maps_list[i] = self.upsample_mask(uncertainty_maps_list[i], mask.detach()) else: for i in range(len(pred_depths_r_list)): pred_depths_r_list[i] = upsample(pred_depths_r_list[i], scale_factor=4) for i in range(len(pred_depths_c_list)): pred_depths_c_list[i] = upsample(pred_depths_c_list[i], scale_factor=4) for i in range(len(uncertainty_maps_list)): uncertainty_maps_list[i] = upsample(uncertainty_maps_list[i], scale_factor=4) return pred_depths_r_list, pred_depths_c_list, uncertainty_maps_list class DispHead(nn.Module): def __init__(self, input_dim=100): super(DispHead, self).__init__() # self.norm1 = nn.BatchNorm2d(input_dim) self.conv1 = nn.Conv2d(input_dim, 1, 3, padding=1) # self.relu = nn.ReLU(inplace=True) self.sigmoid = nn.Sigmoid() def forward(self, x, scale): # x = self.relu(self.norm1(x)) x = self.sigmoid(self.conv1(x)) if scale > 1: x = upsample(x, scale_factor=scale) return x class BasicUpdateBlockDepth(nn.Module): def __init__(self, hidden_dim=128, context_dim=192): super(BasicUpdateBlockDepth, self).__init__() self.encoder = ProjectionInputDepth(hidden_dim=hidden_dim, out_chs=hidden_dim * 2) self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=self.encoder.out_chs+context_dim) self.p_head = PHead(hidden_dim, hidden_dim) def forward(self, depth, context, gru_hidden, seq_len, depth_num, min_depth, max_depth): pred_depths_r_list = [] pred_depths_c_list = [] uncertainty_maps_list = [] b, _, h, w = depth.shape depth_range = max_depth - min_depth interval = depth_range / depth_num interval = interval * torch.ones_like(depth) interval = interval.repeat(1, depth_num, 1, 1) interval = torch.cat([torch.ones_like(depth) * min_depth, interval], 1) bin_edges = torch.cumsum(interval, 1) current_depths = 0.5 * (bin_edges[:, :-1] + bin_edges[:, 1:]) index_iter = 0 for i in range(seq_len): input_features = self.encoder(current_depths.detach()) input_c = torch.cat([input_features, context], dim=1) gru_hidden = self.gru(gru_hidden, input_c) pred_prob = self.p_head(gru_hidden) depth_r = (pred_prob * current_depths.detach()).sum(1, keepdim=True) pred_depths_r_list.append(depth_r) uncertainty_map = torch.sqrt((pred_prob * ((current_depths.detach() - depth_r.repeat(1, depth_num, 1, 1))**2)).sum(1, keepdim=True)) uncertainty_maps_list.append(uncertainty_map) index_iter = index_iter + 1 pred_label = get_label(torch.squeeze(depth_r, 1), bin_edges, depth_num).unsqueeze(1) depth_c = torch.gather(current_depths.detach(), 1, pred_label.detach()) pred_depths_c_list.append(depth_c) label_target_bin_left = pred_label target_bin_left = torch.gather(bin_edges, 1, label_target_bin_left) label_target_bin_right = (pred_label.float() + 1).long() target_bin_right = torch.gather(bin_edges, 1, label_target_bin_right) bin_edges, current_depths = update_sample(bin_edges, target_bin_left, target_bin_right, depth_r.detach(), pred_label.detach(), depth_num, min_depth, max_depth, uncertainty_map) return pred_depths_r_list, pred_depths_c_list, uncertainty_maps_list class PHead(nn.Module): def __init__(self, input_dim=128, hidden_dim=128): super(PHead, self).__init__() self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1) self.conv2 = nn.Conv2d(hidden_dim, 16, 3, padding=1) def forward(self, x): out = torch.softmax(self.conv2(F.relu(self.conv1(x))), 1) return out class SepConvGRU(nn.Module): def __init__(self, hidden_dim=128, input_dim=128+192): super(SepConvGRU, self).__init__() self.convz1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2)) self.convr1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2)) self.convq1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2)) self.convz2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0)) self.convr2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0)) self.convq2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0)) def forward(self, h, x): # horizontal hx = torch.cat([h, x], dim=1) z = torch.sigmoid(self.convz1(hx)) r = torch.sigmoid(self.convr1(hx)) q = torch.tanh(self.convq1(torch.cat([r*h, x], dim=1))) h = (1-z) * h + z * q # vertical hx = torch.cat([h, x], dim=1) z = torch.sigmoid(self.convz2(hx)) r = torch.sigmoid(self.convr2(hx)) q = torch.tanh(self.convq2(torch.cat([r*h, x], dim=1))) h = (1-z) * h + z * q return h class ProjectionInputDepth(nn.Module): def __init__(self, hidden_dim, out_chs): super().__init__() self.out_chs = out_chs self.convd1 = nn.Conv2d(16, hidden_dim, 7, padding=3) self.convd2 = nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1) self.convd3 = nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1) self.convd4 = nn.Conv2d(hidden_dim, out_chs, 3, padding=1) def forward(self, depth): d = F.relu(self.convd1(depth)) d = F.relu(self.convd2(d)) d = F.relu(self.convd3(d)) d = F.relu(self.convd4(d)) return d class Projection(nn.Module): def __init__(self, in_chs, out_chs): super().__init__() self.conv = nn.Conv2d(in_chs, out_chs, 3, padding=1) def forward(self, x): out = self.conv(x) return out def upsample(x, scale_factor=2, mode="bilinear", align_corners=False): """Upsample input tensor by a factor of 2 """ return F.interpolate(x, scale_factor=scale_factor, mode=mode, align_corners=align_corners) def upsample1(x, scale_factor=2, mode="bilinear"): """Upsample input tensor by a factor of 2 """ return F.interpolate(x, scale_factor=scale_factor, mode=mode)