Method forward

model/splat_model.py:169–258

(self, data, step_ratio=0.0)

    def forward(self, data, step_ratio=0.0):
        # data: [B, V, C, H, W]
        input_frames = data['frames']  # [B, V, C, H, W], input features
        input_depths = data['depths']  # [B, V, C, H, W], input features
        timestamps = data['timestamps']  # [B, V], input timestamps
        supv_masks = data['supv_masks']  # [B, V], input timestamps
        if supv_masks is None:
            supv_masks = torch.ones_like(input_depths, device=timestamps.device).bool()
        timestamps = torch.as_tensor(timestamps, dtype=torch.float32, device=input_frames.device)
        timestamps = timestamps / timestamps[..., -1].unsqueeze(-1)
        anchor_time = torch.tensor([0.0, 1.0], device=input_frames.device)
        supv_masks = (supv_masks > 0)
       
        max_depth = input_depths.flatten(2).max(dim=2)[0][:, :, None, None, None]
        min_depth = input_depths.flatten(2).min(dim=2)[0][:, :, None, None, None]
        target_depth = input_depths # without normalization
        input_depths = (input_depths - min_depth) / (max_depth - min_depth)
        
        results = {}
        
        # predict gaussians
        decoder_out = self.forward_gaussians(input_frames, input_depths, timestamps)  # dict
        # pdb.set_trace()

        with autocast('cuda', enabled=False):
            render_pkg = self.gaussian_renderer(decoder_out["pred_gs"], self.background, 
                                                opt=self.opt, timestamps=timestamps, 
                                                anchor_time=anchor_time,
                                                training=self.training,
                                                )
            
            loss, mse_loss, supv_mse_loss, depth_loss, loss_lpips, psnr, metrics = self.compute_losses(input_frames, target_depth, supv_masks, render_pkg)

            if hasattr(self.opt, 'fix_opacity') and self.opt.fix_opacity:
                render_pkg_fix = self.gaussian_renderer(decoder_out["pred_gs"], self.background, 
                                                    opt=self.opt, timestamps=timestamps,
                                                    anchor_time=anchor_time,
                                                    override_opacity=True, training=self.training,
                                                    )
                loss_fix, mse_loss_fix, supv_mse_loss_fix, depth_loss_fix, loss_lpips_fix, psnr_fix, metrics_fix = self.compute_losses(input_frames, target_depth, supv_masks, render_pkg_fix)
                loss = (loss + loss_fix) * 0.5
            else:
                render_pkg_fix = None
                mse_loss_fix = torch.zeros_like(mse_loss)
                depth_loss_fix = torch.zeros_like(depth_loss)  
                loss_lpips_fix = torch.zeros_like(loss_lpips)
                supv_mse_loss_fix = torch.zeros_like(supv_mse_loss)
                psnr_fix = torch.zeros_like(psnr)
        
        pred_depth = render_pkg["depth"]
        depth_mask = pred_depth >= 0.2  # [B, V, 1, H, W]
        pred_depth[~depth_mask] += 10  # set invalid depth to max depth
        pred_depth = 1.0 / (pred_depth + 1e-8)
        B, V, C, H, W = pred_depth.shape
        # Reshape to treat each depth map independently for min/max calculation
        reshaped_depth = pred_depth.view(B, V*C * H * W) # Shape [B*V, H*W]
        min_vals = reshaped_depth.min(dim=1, keepdim=True)[0] # Shape [B*V, 1]
        max_vals = reshaped_depth.max(dim=1, keepdim=True)[0] # Shape [B*V, 1]