深度监督模块移植流程-优快云博客

本文链接：https://blog.youkuaiyun.com/Ajaxm/article/details/146523816

一、整体流程概述

Sat-nerf原本是一个基于NeRF的卫星图像三维重建模型，它主要依赖颜色一致性来优化三维几何结构。我们的修改目标是引入深度监督，让模型能够学习更准确的几何信息。整个流程包括以下步骤：

从单目深度估计模型获取相对深度
将相对深度与稀疏点云对齐，获取绝对深度
使用这些深度作为额外的监督信号训练网络
预计算和缓存深度以提高效率

二、单目深度估计与深度对齐

1. 深度模型初始化

首先，在SatelliteDataset_depth类中添加了深度模型的延迟加载功能：

def _load_depth_model(self):
    # 仅在需要时才加载模型
    if not hasattr(self, 'depth_model') or self.depth_model is None:
        logger.info("Loading Depth Anything V2 model...")
        # 选择模型大小（小、中、大）
        encoder = 'vitb'  
        # 加载模型配置和权重
        self.depth_model = DepthAnythingV2(**model_configs[encoder])
        self.depth_model.load_state_dict(torch.load(checkpoint_path))
        self.depth_model = self.depth_model.to('cuda').eval()

这一步的目的是引入预训练的Depth Anything V2模型，它可以从单视图卫星图像中预测出相对深度图。

2. 深度尺度优化

单目深度估计模型只能提供相对深度，我们需要将其转换为绝对深度。这通过与稀疏点云（来自bundle adjustment）对齐实现：

def optimize_depth(self, sparse_depths, rel_depths, weights=None):
    # 确保输入是张量
    sparse_depths = sparse_depths.to(device)
    rel_depths = rel_depths.to(device)
    
    # 初始化可优化参数
    scale = torch.tensor([1.0], requires_grad=True, device=device)
    offset = torch.tensor([0.0], requires_grad=True, device=device)
    
    # 通过梯度下降找到最佳的尺度和偏移
    optimizer = torch.optim.Adam([scale, offset], lr=0.01)
    
    for i in range(300):
        # 计算预测深度
        pred_depths = scale * rel_depths + offset
        
        # 计算加权MSE损失
        loss = torch.mean(weights_valid * (pred_depths - sparse_valid) ** 2)
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    return best_scale, best_offset

这个过程的核心是公式：D_dense = s·D̂_dense + o，其中s和o是通过最小化稀疏点处的深度差异得到的缩放和偏移参数。这与Sat-DN论文中的公式6和7对应。

三、深度缓存与预计算

为了加速训练，我们开发了深度预计算系统：

1. 预计算脚本

在satdp_precompute_depths.py中实现了完整的预计算流程：

def process_dataset(args):
    # 加载深度模型
    depth_model = load_depth_model(model_size=args.model_size)
    
    # 遍历所有训练图像
    for json_p in json_files:
        # 读取图像和相机信息
        d = sat_utils.read_dict_from_json(json_p)
        img_id = sat_utils.get_file_id(d["img"])
        img_p = os.path.join(img_dir, d["img"])
        
        # 获取稀疏深度和点信息
        pts2d = np.array(d["keypoints"]["2d_coordinates"]) / args.img_downscale
        pts3d_indices = d["keypoints"]["pts3d_indices"]
        pts3d = np.array(tie_points[pts3d_indices, :])
        
        # 计算稀疏深度（根据RPC模型）
        # 这里对应Sat-DN论文中的公式5
        sparse_depths = torch.linalg.norm(pts3d_tensor - rays_o_tensor, dim=1)
        
        # 预测单目深度
        dense_depth = depth_model.infer_image(img)
        
        # 提取稀疏点位置的相对深度
        rel_depths = [深度图在每个点位置的值]
        
        # 优化尺度和偏移
        scale, offset = optimize_depth(sparse_depths, rel_depths, point_weights)
        
        # 保存到缓存
        np.savez(cache_path, scale=scale, offset=offset, ...)

这个脚本可以提前处理所有训练图像的深度信息，避免每次训练时重复计算。

2. 缓存加载逻辑

在训练时，修改了load_depth_data函数来优先使用缓存：

def load_depth_data(self, json_files, tie_points, verbose=False):
    # 检查是否有缓存可用
    depth_cache_dir = os.path.join(self.cache_dir, "depth_cache")
    use_cache = os.path.exists(depth_cache_dir)
    
    if use_cache:
        logger.info(f"使用预计算的深度缓存")
    else:
        logger.info("未找到深度缓存，将实时计算")
        self._load_depth_model()
    
    # 处理每个图像
    for t, json_p in enumerate(json_files):
        # ... 获取图像信息 ...
        
        # 根据是否有缓存决定加载方式
        if use_cache:
            cache_path = os.path.join(depth_cache_dir, f"{img_id}_depth.npz")
            if os.path.exists(cache_path):
                cache_data = np.load(cache_path)
                cal_depths = torch.from_numpy(cache_data['calibrated_depths']).float()
            else:
                # 缓存不存在时使用原始深度
                cal_depths = sparse_depths
        else:
            # 实时计算深度
            dense_depth = self.depth_model.infer_image(img)
            rel_depths = [从dense_depth中提取点深度]
            scale, offset = self.optimize_depth(sparse_depths, rel_depths, point_weights)
            cal_depths = scale * rel_depths + offset

这个优化大大提高了训练效率，避免了每次都需要加载深度模型和计算深度对齐。

四、深度损失的引入

1. 修改损失函数

在satdp_metrics.py中，将原本的MSE深度损失修改为L1损失：

class DepthLoss(torch.nn.Module):
    def __init__(self, lambda_ds=1.0):
        super().__init__()
        self.lambda_ds = lambda_ds/3.
        # 使用L1损失替代MSE损失，与Sat-DN论文公式9一致
        self.loss = torch.nn.L1Loss(reduction='none')

    def forward(self, inputs, targets, weights=1.):
        # 计算深度损失
        loss_dict = {}
        typ = 'coarse'
        loss_dict[f'{typ}_ds'] = self.loss(inputs['depth_coarse'], targets)
        
        # 应用权重（从点云重投影误差获得的可靠性权重）
        for k in loss_dict.keys():
            loss_dict[k] = self.lambda_ds * torch.mean(weights * loss_dict[k])
        
        return sum(loss_dict.values()), loss_dict

L1损失比MSE损失对深度异常值更健壮，这与Sat-DN论文中的公式9一致。

2. 渲染深度

为了与深度监督信号比较，需要从模型中渲染深度：

# 与颜色渲染类似，深度渲染也使用体积渲染公式
def render_depth(ray):
    D(r) = Σ T_i·α_i·d_i

这个过程复用了颜色渲染时计算的权重，不需要额外的计算开销。

五、流程对比与总结

原始Sat-nerf流程:

从卫星图像和RPC参数计算光线
采样光线上的点
通过MLP预测每个点的密度和颜色
通过体积渲染合成图像
比较渲染图像与真实图像，优化MLP权重

修改后的流程:

从卫星图像和RPC参数计算光线
预计算/加载图像的深度图（新增）
采样光线上的点
通过MLP预测每个点的密度和颜色
通过体积渲染合成图像和深度图
比较渲染图像与真实图像，同时比较渲染深度与预计算深度（新增）
优化MLP权重

关键区别:

额外的监督信号: 原始Sat-nerf仅使用颜色监督，修改后添加了深度监督
几何先验: 通过预训练的单目深度模型引入了几何先验知识
损失权重平衡: 需要平衡颜色损失与深度损失的权重
预计算与缓存: 添加了深度预计算和缓存机制提高效率

通过这些修改，模型能够学习到更准确的几何信息，特别是在纹理缺乏或光照变化较大的区域。这与Sat-DN论文中的深度监督策略一致，成功地将Depth-Regularized论文中的思想移植到了Sat-nerf中。