Method compute_losses

model/splat_model.py:69–166

(self, input_frames, target_depth, supv_masks, render_pkg)

        return decoder_out
    
    def compute_losses(self, input_frames, target_depth, supv_masks, render_pkg):
        output_frames = render_pkg["render"]  # [B, V, C, H, W]
        pred_depths = render_pkg["depth"]
        depth_mask = render_pkg["alpha"] > 0.1  # [B, V, 1, H, W]
        metrics = {}
        with torch.no_grad():
            B, V, C, H, W = input_frames.shape
            
            # All frames
            input_frames_all_256 = F.interpolate(input_frames.reshape(-1, 3, H, W), (256, 256), mode='bilinear', align_corners=False)
            output_frames_all_256 = F.interpolate(output_frames.reshape(-1, 3, H, W), (256, 256), mode='bilinear', align_corners=False)
            metrics['psnr'] = compute_psnr(input_frames_all_256, output_frames_all_256).mean()
            metrics['ssim'] = compute_ssim(input_frames_all_256, output_frames_all_256).mean()
            metrics['lpips'] = compute_lpips(input_frames_all_256 * 2 - 1, output_frames_all_256 * 2 - 1).mean()

            metrics['input_frames'] = input_frames
            metrics['pred_frames'] = output_frames

            # Middle frames [1:-1]
            if V > 2:
                input_frames_middle = input_frames[:, 1:-1]
                output_frames_middle = output_frames[:, 1:-1]
                
                input_frames_middle_256 = F.interpolate(input_frames_middle.reshape(-1, 3, H, W), (256, 256), mode='bilinear', align_corners=False)
                output_frames_middle_256 = F.interpolate(output_frames_middle.reshape(-1, 3, H, W), (256, 256), mode='bilinear', align_corners=False)

                metrics['psnr_novel'] = compute_psnr(input_frames_middle_256, output_frames_middle_256).mean()
                metrics['ssim_novel'] = compute_ssim(input_frames_middle_256, output_frames_middle_256).mean()
                metrics['lpips_novel'] = compute_lpips(input_frames_middle_256 * 2 - 1, output_frames_middle_256 * 2 - 1).mean()
            else:
                metrics['psnr_novel'] = torch.tensor(0.0, device=input_frames.device)
                metrics['ssim_novel'] = torch.tensor(0.0, device=input_frames.device)
                metrics['lpips_novel'] = torch.tensor(0.0, device=input_frames.device)

            # First and last frames (given views)
            if V >= 1:
                indices = [0]
                if V > 1:
                    indices.append(V - 1)
                
                input_frames_ends = input_frames[:, indices]
                output_frames_ends = output_frames[:, indices]

                input_frames_ends_256 = F.interpolate(input_frames_ends.reshape(-1, 3, H, W), (256, 256), mode='bilinear', align_corners=False)
                output_frames_ends_256 = F.interpolate(output_frames_ends.reshape(-1, 3, H, W), (256, 256), mode='bilinear', align_corners=False)
                
                metrics['psnr_given'] = compute_psnr(input_frames_ends_256, output_frames_ends_256).mean()
                metrics['ssim_given'] = compute_ssim(input_frames_ends_256, output_frames_ends_256).mean()
                metrics['lpips_given'] = compute_lpips(input_frames_ends_256 * 2 - 1, output_frames_ends_256 * 2 - 1).mean()
            else: # Should not happen if V >= 1
                metrics['psnr_given'] = torch.tensor(0.0, device=input_frames.device)
                metrics['ssim_given'] = torch.tensor(0.0, device=input_frames.device)
                metrics['lpips_given'] = torch.tensor(0.0, device=input_frames.device)
            
            psnr = metrics['psnr']

        if self.opt.skip:
            # skip the frame_0