model/compute_graph/graph_shape.py

import torch
import torch.nn as nn

from utils.util import EasyDict as edict
from utils.loss import Loss
from model.shape.implicit import Implicit
from model.shape.seen_coord_enc import CoordEncAtt, CoordEncRes
from model.shape.rgb_enc import RGBEncAtt, RGBEncRes
from model.depth.dpt_depth import DPTDepthModel
from utils.util import toggle_grad, interpolate_coordmap, get_child_state_dict
from utils.camera import unproj_depth, valid_norm_fac
from utils.layers import Bottleneck_Conv

class Graph(nn.Module):

    def __init__(self, opt):
        super().__init__()
        # define the intrinsics head
        self.intr_feat_channels = 768
        self.intr_head = nn.Sequential(
            Bottleneck_Conv(self.intr_feat_channels, kernel_size=3),
            Bottleneck_Conv(self.intr_feat_channels, kernel_size=3),
        )
        self.intr_pool = nn.AdaptiveAvgPool2d((1, 1))
        self.intr_proj = nn.Linear(self.intr_feat_channels, 3)
        # init the last linear layer so it outputs zeros
        nn.init.zeros_(self.intr_proj.weight)
        nn.init.zeros_(self.intr_proj.bias)
            
        # define the depth pred model based on omnidata
        self.dpt_depth = DPTDepthModel(backbone='vitb_rn50_384')
        # load the pretrained depth model
        # when intrinsics need to be predicted we need to load that part as well
        self.load_pretrained_depth(opt)
        if opt.optim.fix_dpt:
            toggle_grad(self.dpt_depth, False)
            toggle_grad(self.intr_head, False)
            toggle_grad(self.intr_proj, False)

        # encoder that encode seen surface to impl conditioning vec
        if opt.arch.depth.encoder == 'resnet':
            opt.arch.depth.dsp = 1
            self.coord_encoder = CoordEncRes(opt)
        else:
            self.coord_encoder = CoordEncAtt(embed_dim=opt.arch.latent_dim, n_blocks=opt.arch.depth.n_blocks, 
                                        num_heads=opt.arch.num_heads, win_size=opt.arch.win_size//opt.arch.depth.dsp)
            
        # rgb branch (not used in final model, keep here for extension)
        if opt.arch.rgb.encoder == 'resnet':
            self.rgb_encoder = RGBEncRes(opt)
        elif opt.arch.rgb.encoder == 'transformer':
            self.rgb_encoder = RGBEncAtt(img_size=opt.H, embed_dim=opt.arch.latent_dim, n_blocks=opt.arch.rgb.n_blocks, 
                                        num_heads=opt.arch.num_heads, win_size=opt.arch.win_size)
        else:
            self.rgb_encoder = None
        
        # implicit function
        feat_res = opt.H // opt.arch.win_size
        self.impl_network = Implicit(feat_res**2, latent_dim=opt.arch.latent_dim*2 if self.rgb_encoder else opt.arch.latent_dim, 
                                     semantic=self.rgb_encoder is not None, n_channels=opt.arch.impl.n_channels, 
                                     n_blocks_attn=opt.arch.impl.att_blocks, n_layers_mlp=opt.arch.impl.mlp_layers, 
                                     num_heads=opt.arch.num_heads, posenc_3D=opt.arch.impl.posenc_3D, 
                                     mlp_ratio=opt.arch.impl.mlp_ratio, skip_in=opt.arch.impl.skip_in, 
                                     pos_perlayer=opt.arch.impl.posenc_perlayer)
        
        # loss functions
        self.loss_fns = Loss(opt)
            
    def load_pretrained_depth(self, opt):
        if opt.pretrain.depth:
            # loading from our pretrained depth and intr model
            if opt.device == 0:
                print("loading dpt depth from {}...".format(opt.pretrain.depth))
            checkpoint = torch.load(opt.pretrain.depth, map_location="cuda:{}".format(opt.device))
            self.dpt_depth.load_state_dict(get_child_state_dict(checkpoint["graph"], "dpt_depth"))
            # load the intr head
            if opt.device == 0:
                print("loading pretrained intr from {}...".format(opt.pretrain.depth))
            self.intr_head.load_state_dict(get_child_state_dict(checkpoint["graph"], "intr_head"))
            self.intr_proj.load_state_dict(get_child_state_dict(checkpoint["graph"], "intr_proj"))
        elif opt.arch.depth.pretrained:
            # loading from omnidata weights
            if opt.device == 0:
                print("loading dpt depth from {}...".format(opt.arch.depth.pretrained))
            checkpoint = torch.load(opt.arch.depth.pretrained, map_location="cuda:{}".format(opt.device))
            state_dict = checkpoint['model_state_dict']
            self.dpt_depth.load_state_dict(state_dict)

    def intr_param2mtx(self, opt, intr_params):
        '''
        Parameters:
            opt: config
            intr_params: [B, 3], [scale_f, delta_cx, delta_cy]
        Return:
            intr: [B, 3, 3]
        '''
        batch_size = len(intr_params)
        f = 1.3875
        intr = torch.zeros(3, 3).float().to(intr_params.device).unsqueeze(0).repeat(batch_size, 1, 1)
        intr[:, 2, 2] += 1
        # scale the focal length
        # range: [-1, 1], symmetric
        scale_f = torch.tanh(intr_params[:, 0])
        # range: [1/4, 4], symmetric
        scale_f = torch.pow(4. , scale_f)
        intr[:, 0, 0] += f * opt.W * scale_f
        intr[:, 1, 1] += f * opt.H * scale_f
        # shift the optic center, (at most to the image border)
        shift_cx = torch.tanh(intr_params[:, 1]) * opt.W / 2
        shift_cy = torch.tanh(intr_params[:, 2]) * opt.H / 2
        intr[:, 0, 2] += opt.W / 2 + shift_cx
        intr[:, 1, 2] += opt.H / 2 + shift_cy
        return intr

    def forward(self, opt, var, training=False, get_loss=True):
        batch_size = len(var.idx)
        
        # encode the rgb, [B, 3, H, W] -> [B, 1+H/(ws)*W/(ws), C], not used in our final model
        var.latent_semantic = self.rgb_encoder(var.rgb_input_map) if self.rgb_encoder else None

        # predict the depth map and intrinsics
        var.depth_pred, intr_feat = self.dpt_depth(var.rgb_input_map, get_feat=True)
        depth_map = var.depth_pred
        # predict the intrinsics
        intr_feat = self.intr_head(intr_feat)
        intr_feat = self.intr_pool(intr_feat).squeeze(-1).squeeze(-1)
        intr_params = self.intr_proj(intr_feat)
        # [B, 3, 3]
        var.intr_pred = self.intr_param2mtx(opt, intr_params)
        intr_forward = var.intr_pred
        # record the validity mask, [B, H*W]
        var.validity_mask = (var.mask_input_map>0.5).float().view(batch_size, -1)

        # project the depth to 3D points in view-centric frame
        # [B, H*W, 3], in camera coordinates
        seen_points_3D_pred = unproj_depth(opt, depth_map, intr_forward)
        # [B, H*W, 3], [B, 1, H, W] (boolean) -> [B, 3], [B]
        seen_points_mean_pred, seen_points_scale_pred = valid_norm_fac(seen_points_3D_pred, var.mask_input_map > 0.5)
        # normalize the seen surface, [B, H*W, 3]
        var.seen_points = (seen_points_3D_pred - seen_points_mean_pred.unsqueeze(1)) / seen_points_scale_pred.unsqueeze(-1).unsqueeze(-1)
        var.seen_points[(var.mask_input_map<=0.5).view(batch_size, -1)] = 0
        # [B, 3, H, W]
        seen_3D_map = var.seen_points.view(batch_size, opt.H, opt.W, 3).permute(0, 3, 1, 2).contiguous()
        seen_3D_dsp, mask_dsp = interpolate_coordmap(seen_3D_map, var.mask_input_map, (opt.H//opt.arch.depth.dsp, opt.W//opt.arch.depth.dsp))
        
        # encode the depth, [B, 1, H/k, W/k] -> [B, 1+H/(ws)*W/(ws), C]
        if opt.arch.depth.encoder == 'resnet':
            var.latent_depth = self.coord_encoder(seen_3D_dsp, mask_dsp)
        else:
            var.latent_depth = self.coord_encoder(seen_3D_dsp.permute(0, 2, 3, 1).contiguous(), mask_dsp.squeeze(1)>0.5)
        

        var.pose = var.pose_gt
        # forward for loss calculation (only during training)
        if 'gt_sample_points' in var and 'gt_sample_sdf' in var:
            with torch.no_grad():
                # get the normalizing fac based on the GT seen surface
                # project the GT depth to 3D points in view-centric frame
                # [B, H*W, 3], in camera coordinates
                seen_points_3D_gt = unproj_depth(opt, var.depth_input_map, var.intr)
                # [B, H*W, 3], [B, 1, H, W] (boolean) -> [B, 3], [B]
                seen_points_mean_gt, seen_points_scale_gt = valid_norm_fac(seen_points_3D_gt, var.mask_input_map > 0.5)
                var.seen_points_gt = (seen_points_3D_gt - seen_points_mean_gt.unsqueeze(1)) / seen_points_scale_gt.unsqueeze(-1).unsqueeze(-1)
                var.seen_points_gt[(var.mask_input_map<=0.5).view(batch_size, -1)] = 0
                
                # transform the GT points accordingly
                # [B, 3, 3]
                R_gt = var.pose_gt[:, :, :3]
                # [B, 3, 1]
                T_gt = var.pose_gt[:, :, 3:]
                # [B, 3, N]
                gt_sample_points_transposed = var.gt_sample_points.permute(0, 2, 1).contiguous()
                # camera coordinates, [B, N, 3]
                gt_sample_points_cam = (R_gt @ gt_sample_points_transposed + T_gt).permute(0, 2, 1).contiguous()
                # normalize with seen std and mean, [B, N, 3]
                var.gt_points_cam = (gt_sample_points_cam - seen_points_mean_gt.unsqueeze(1)) / seen_points_scale_gt.unsqueeze(-1).unsqueeze(-1)
                
                # get near-surface points for visualization
                # [B, 100, 3]
                close_surf_idx = torch.topk(var.gt_sample_sdf.abs(), k=100, dim=1, largest=False)[1].unsqueeze(-1).repeat(1, 1, 3)
                # [B, 100, 3]
                var.gt_surf_points = torch.gather(var.gt_points_cam, dim=1, index=close_surf_idx)
        
            # [B, N], [B, N, 1+feat_res**2], inference the impl_network for 3D loss
            var.pred_sample_occ, attn = self.impl_network(var.latent_depth, var.latent_semantic, var.gt_points_cam)

        # calculate the loss if needed
        if get_loss: 
            loss = self.compute_loss(opt, var, training)
            return var, loss
        
        return var

    def compute_loss(self, opt, var, training=False):
        loss = edict()
        if opt.loss_weight.depth is not None:
            loss.depth = self.loss_fns.depth_loss(var.depth_pred, var.depth_input_map, var.mask_input_map)
        if opt.loss_weight.intr is not None and training:
            loss.intr = self.loss_fns.intr_loss(var.seen_points, var.seen_points_gt, var.validity_mask)
        if opt.loss_weight.shape is not None and training:
            loss.shape = self.loss_fns.shape_loss(var.pred_sample_occ, var.gt_sample_sdf)
        return loss