SurroundOcc 代码的 数据加载，推理，指标评测，可视化

Occupancy prediction:

demo

Occupancy prediction:demo

简介

为了实现更全面、更一致的场景重建，本文提出了一种 SurroundOcc 方法来预测多摄像头图像的体积占位。我们首先提取每幅图像的多尺度特征，并采用空间交叉注意将其提升到三维体积空间。然后，我们应用三维卷积来逐步提高体积特征的采样率，并施加多层次的监督。为了训练多摄像头三维场景重建模型，我们设计了一个管道，利用稀疏的激光雷达点生成密集的占地地面实况。生成管道只需要存在的三维检测和三维语义分割标签，无需额外的人工注释。具体来说，我们分别对动态物体和静态场景的多帧 LiDAR 点进行融合。然后，我们采用泊松重构（Poisson Reconstruction）来填补漏洞，并对网格进行体素化处理，从而获得密集的体积占位。

数据加载

https://github.com/weiyithu/SurroundOcc/blob/main/projects/mmdet3d_plugin/datasets/pipelines/loading.py

class LoadOccupancy(object):
    """
    加载占据空间的地面真实数据。

    期望results['occ_path']是文件名列表。

    地面真实数据是一个(N, 4)的张量，N是被占据体素的数量，
    前三个通道代表xyz体素坐标，最后一个通道是语义类别。
    """

    def __init__(self, use_semantic=True):
        """
        初始化函数。
        
        Args:
            use_semantic (bool): 是否使用语义信息，默认为True。
        """
        self.use_semantic = use_semantic

    def __call__(self, results):
        """
        使实例可以像函数那样被调用。

        Args:
            results (dict): 包含加载信息的字典，期望有'occ_path'键。

        Returns:
            dict: 更新后包含地面真实占据空间数据的results字典。
        """
        # 根据路径加载占据空间数据并转换为float32类型
        occ = np.load(results['occ_path'])
        occ = occ.astype(np.float32)
        
        # 如果使用语义信息
        if self.use_semantic:
            # 将语义类别为0的体素类别标记为255（忽略类）
            occ[..., 3][occ[..., 3] == 0] = 255
        else:
            # 如果不使用语义信息，只保留类别大于0的体素，并将它们的类别设置为1
            occ = occ[occ[..., 3] > 0]
            occ[..., 3] = 1
        
        # 将处理后的占据空间数据保存回results字典
        results['gt_occ'] = occ
        
        return results

    def __repr__(self):
        """
        返回描述该模块的字符串。

        Returns:
            str: 描述模块的字符串。
        """
        repr_str = self.__class__.__name__
        return repr_str

推理

https://github.com/weiyithu/SurroundOcc/blob/main/projects/configs/surroundocc/surroundocc_inference.py

配置文件，用于设置任务中的数据集、模型参数、训练和测试管道等

# 基础配置文件，指定了数据集和运行时的基本设置
_base_ = [
    '../datasets/custom_nus-3d.py',  # 自定义NuScenes数据集的配置
    '../_base_/default_runtime.py'  # 默认的运行时配置
]

# 插件设置，用于扩展功能
plugin = True
plugin_dir = 'projects/mmdet3d_plugin/'

# 点云的范围，改变点云范围时，模型应相应地改变其点云范围
point_cloud_range = [-50, -50, -5.0, 50, 50, 3.0]
occ_size = [200, 200, 16]  # 占据空间的大小
use_semantic = True  # 是否使用语义信息

# 图像归一化配置
img_norm_cfg = dict(
    mean=[103.530, 116.280, 123.675],  # 平均值
    std=[1.0, 1.0, 1.0],  # 标准差
    to_rgb=False)  # 是否转换为RGB格式

# 类别名称
class_names = ['barrier', 'bicycle', 'bus', 'car', 'construction_vehicle', 'motorcycle',
               'pedestrian', 'traffic_cone', 'trailer', 'truck', 'driveable_surface',
               'other_flat', 'sidewalk', 'terrain', 'manmade', 'vegetation']

# 输入模态配置
input_modality = dict(
    use_lidar=False,  # 是否使用激光雷达数据
    use_camera=True,  # 是否使用相机数据
    use_radar=False,  # 是否使用雷达数据
    use_map=False,  # 是否使用地图数据
    use_external=True)  # 是否使用外部数据

# 模型和数据集的其他配置参数
_dim_ = [128, 256, 512]
_ffn_dim_ = [256, 512, 1024]
volume_h_ = [100, 50, 25]
volume_w_ = [100, 50, 25]
volume_z_ = [8, 4, 2]
_num_points_ = [2, 4, 8]
_num_layers_ = [1, 3, 6]

model = dict(
    # 定义模型的具体参数
)

dataset_type = 'CustomNuScenesOccDataset'
data_root = 'data/nuscenes/'
file_client_args = dict(backend='disk')

train_pipeline = [
    # 训练数据处理管道
]

test_pipeline = [
    # 测试数据处理管道
]

find_unused_parameters = True
data = dict(
    # 数据配置
)

optimizer = dict(
    weight_decay=0.01)  # 优化器配置，包括权重衰减等

optimizer_config = dict(grad_clip=dict(max_norm=35, norm_type=2))
# 学习策略配置

lr_config = dict(
    min_lr_ratio=1e-3)
    
total_epochs = 24  # 总训练轮次

evaluation = dict(interval=1, pipeline=test_pipeline)  # 评估配置

runner = dict(type='EpochBasedRunner', max_epochs=total_epochs)
load_from = 'ckpts/r101_dcn_fcos3d_pretrain.pth'  # 预训练模型
log_config = dict(
    # 日志配置
)

checkpoint_config = dict(interval=1)  # 检查点配置

指标评测

def voxel_to_vertices(voxel, img_metas, thresh=0.5):
    """
    将体素数据转换为顶点集。
    
    参数:
    - voxel: 体素数据，一个三维张量。
    - img_metas: 包含图像元数据的字典。
    - thresh: 阈值，用于确定体素是否被视为占用。
    
    返回:
    - vertices: 根据体素数据生成的顶点集，每个顶点代表一个体素的中心坐标。
    """
    # 创建三维网格，对应于体素的每个位置
    x = torch.linspace(0, voxel.shape[0] - 1, voxel.shape[0])
    y = torch.linspace(0, voxel.shape[1] - 1, voxel.shape[1])
    z = torch.linspace(0, voxel.shape[2] - 1, voxel.shape[2])
    X, Y, Z = torch.meshgrid(x, y, z)
    vv = torch.stack([X, Y, Z], dim=-1).to(voxel.device)

    # 筛选出阈值以上的体素，将它们的坐标转换为实际空间坐标
    vertices = vv[voxel > thresh]
    for i in range(3):
        vertices[:, i] = (vertices[:, i] + 0.5) * (img_metas['pc_range'][i+3] - img_metas['pc_range'][i]) / img_metas['occ_size'][i] + img_metas['pc_range'][i]

    return vertices

def gt_to_vertices(gt, img_metas):
    """
    将地面真实标注转换为顶点集。
    
    参数:
    - gt: 地面真实标注数据，每行代表一个顶点的xyz坐标和类别。
    - img_metas: 包含图像元数据的字典。
    
    返回:
    - gt: 转换后的顶点集。
    """
    for i in range(3):
        gt[:, i] = (gt[:, i] + 0.5) * (img_metas['pc_range'][i+3] - img_metas['pc_range'][i]) / img_metas['occ_size'][i] + img_metas['pc_range'][i]
    return gt

def gt_to_voxel(gt, img_metas):
    """
    将地面真实标注转换为体素表示。
    
    参数:
    - gt: 地面真实标注数据。
    - img_metas: 包含图像元数据的字典。
    
    返回:
    - voxel: 体素数据，其中体素值代表语义类别。
    """
    voxel = np.zeros(img_metas['occ_size'])
    voxel[gt[:, 0].astype(np.int), gt[:, 1].astype(np.int), gt[:, 2].astype(np.int)] = gt[:, 3]

    return voxel

def eval_3d(verts_pred, verts_trgt, threshold=.5):
    """
    使用查谟距离评估3D重建的质量。

    参数:
    - verts_pred: 预测的顶点集，一个张量。
    - verts_trgt: 目标（真实）顶点集，一个张量。
    - threshold: 计算精确度和召回率时使用的距离阈值。

    返回:
    - metrics: 包含平均距离、查谟距离、精确度、召回率和F1分数的数组。
    """
    # 计算查谟距离
    d1, d2, idx1, idx2 = chamfer.forward(verts_pred.unsqueeze(0).type(torch.float), verts_trgt.unsqueeze(0).type(torch.float))
    # 计算每个点的距离
    dist1 = torch.sqrt(d1).cpu().numpy()
    dist2 = torch.sqrt(d2).cpu().numpy()
    # 计算查谟距离的平均值
    cd = dist1.mean() + dist2.mean()
    # 计算精确度和召回率
    precision = np.mean((dist1 < threshold).astype('float'))
    recal = np.mean((dist2 < threshold).astype('float'))
    # 计算F1分数
    fscore = 2 * precision * recal / (precision + recal)
    metrics = np.array([np.mean(dist1), np.mean(dist2), cd, precision, recal, fscore])
    return metrics

def evaluation_reconstruction(pred_occ, gt_occ, img_metas):
    """
    评估3D重建质量。

    参数:
    - pred_occ: 预测的占据空间（体素）数据。
    - gt_occ: 真实的占据空间数据。
    - img_metas: 图像的元数据信息。

    返回:
    - 一个包含所有样本评估指标的数组。
    """
    results = []
    for i in range(pred_occ.shape[0]):
        # 转换预测和真实的体素数据为顶点集
        vertices_pred = voxel_to_vertices(pred_occ[i], img_metas, thresh=0.25)  # 为了解决类别不平衡问题设置较低的阈值
        vertices_gt = gt_to_vertices(gt_occ[i][..., :3], img_metas)
        
        # 使用查谟距离评估顶点集
        metrics = eval_3d(vertices_pred.type(torch.double), vertices_gt.type(torch.double))  # 由于查谟距离函数的一个bug，必须转换为double类型
        results.append(metrics)
    return np.stack(results, axis=0)

def evaluation_semantic(pred_occ, gt_occ, img_metas, class_num):
    """
    评估语义分割性能。

    参数:
    - pred_occ: 预测的占据空间（体素）数据。
    - gt_occ: 真实的占据空间数据。
    - img_metas: 图像的元数据信息。
    - class_num: 类别的数量。

    返回:
    - 一个包含所有样本每个类别评估分数的数组。
    """
    results = []
    for i in range(pred_occ.shape[0]):
        # 转换为numpy数组以便处理
        gt_i, pred_i = gt_occ[i].cpu().numpy(), pred_occ[i].cpu().numpy()
        # 将真实数据转换为体素表示
        gt_i = gt_to_voxel(gt_i, img_metas)
        mask = (gt_i != 255)  # 忽略特定类别
        score = np.zeros((class_num, 3))
        for j in range(class_num):
            if j == 0:  # 类别0用于计算几何IoU
                score[j][0] += ((gt_i[mask] != 0) * (pred_i[mask] != 0)).sum()
                score[j][1] += (gt_i[mask] != 0).sum()
                score[j][2] += (pred_i[mask] != 0).sum()
            else:
                # 对于其他类别，计算匹配的体素数量
                score[j][0] += ((gt_i[mask] == j) * (pred_i[mask] == j)).sum()
                score[j][1] += (gt_i[mask] == j).sum()
                score[j][2] += (pred_i[mask] == j).sum()
        results.append(score)
    return np.stack(results, axis=0)

可视化

https://github.com/weiyithu/SurroundOcc/blob/main/tools/visual.py

# 一些安装包

# 初始化颜色数组
colors = np.array(
    # 此处应有颜色数据，省略
).astype(np.uint8)  # 将颜色数组的数据类型转换为无符号8位整型

# 设置Mayavi的选项为离屏渲染模式
#mlab.options.offscreen = True

# 设置体素的大小为0.5
voxel_size = 0.5
# 设置点云范围
pc_range = [-50, -50,  -5, 50, 50, 3]

# 从命令行参数获取可视化路径
visual_path = sys.argv[1]
# 加载视场（Field of View, FOV）体素数据
fov_voxels = np.load(visual_path)

# 过滤掉第四维（通常为强度或其他属性）小于等于0的体素
fov_voxels = fov_voxels[fov_voxels[..., 3] > 0]
# 将体素的XYZ坐标根据体素大小进行缩放，并进行偏移，以匹配点云的实际范围
fov_voxels[:, :3] = (fov_voxels[:, :3] + 0.5) * voxel_size
fov_voxels[:, 0] += pc_range[0]
fov_voxels[:, 1] += pc_range[1]
fov_voxels[:, 2] += pc_range[2]

# 创建一个Mayavi图形窗口，设置尺寸和背景颜色
# figure = mlab.figure(size=(600, 600), bgcolor=(1, 1, 1))
figure = mlab.figure(size=(2560, 1440), bgcolor=(1, 1, 1))

# 使用pdb进行调试
# pdb.set_trace()

# 使用Mayavi的points3d方法绘制3D点云，其中每个点的大小和颜色根据其属性值动态调整
plt_plot_fov = mlab.points3d(
    fov_voxels[:, 0],  # X坐标
    fov_voxels[:, 1],  # Y坐标
    fov_voxels[:, 2],  # Z坐标
    fov_voxels[:, 3],  # 根据第四维的值调整颜色
    colormap="viridis",  # 使用viridis颜色图
    scale_factor=voxel_size - 0.05*voxel_size,  # 设置体素缩放因子
    mode="cube",  # 使用立方体表示每个点
    opacity=1.0,  # 设置不透明度为完全不透明
    vmin=0,  # 设置颜色映射的最小值
    vmax=19,  # 设置颜色映射的最大值
)

# 设置体素的缩放模式为“按向量缩放”
plt_plot_fov.glyph.scale_mode = "scale_by_vector"
# 将颜色映射表（Look-Up Table, LUT）设置为之前定义的颜色
plt_plot_fov.module_manager.scalar_lut_manager.lut.table = colors

# 保存Mayavi渲染的图像到文件
#mlab.savefig('temp/mayavi.png')
# 显示Mayavi渲染窗口
mlab.show()