2. 体素化输入表示与特征提取

任务目标

将原始点云转换为固定分辨率的体素网格,并通过3D CNN提取体素特征,为Occupancy网络提供稠密化的输入表示。本案例聚焦于体素化流程与特征提取主干的最小可复现实现。

核心实现

import torch
import torch.nn as nn
import numpy as np

def point_cloud_to_voxel(points, grid_size=32, max_range=1.0):
    """
    将点云转换为二值体素网格
    points: (N, 3) 归一化到 [-max_range, max_range]
    grid_size: 体素分辨率(各维度相同)
    返回: (1, 1, D, H, W) 二值体素张量
    """
    # 坐标裁剪与归一化到 [0, grid_size-1]
    points = np.clip(points, -max_range, max_range)
    voxel_coords = ((points + max_range) / (2 * max_range) * (grid_size - 1)).astype(np.int32)

    voxel_grid = np.zeros((grid_size, grid_size, grid_size), dtype=np.float32)
    # 去重并填充
    unique_coords = np.unique(voxel_coords, axis=0)
    voxel_grid[unique_coords[:, 0], unique_coords[:, 1], unique_coords[:, 2]] = 1.0

    # 增加 batch 和 channel 维度
    return torch.from_numpy(voxel_grid).unsqueeze(0).unsqueeze(0)

class VoxelFeatureExtractor(nn.Module):
    """轻量3D CNN:从体素网格提取128维特征向量"""
    def __init__(self, input_channels=1, feature_dim=128):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv3d(input_channels, 16, kernel_size=3, padding=1),
            nn.BatchNorm3d(16),
            nn.ReLU(inplace=True),
            nn.MaxPool3d(2),  # 16^3

            nn.Conv3d(16, 32, kernel_size=3, padding=1),
            nn.BatchNorm3d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool3d(2),  # 8^3

            nn.Conv3d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm3d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool3d(2),  # 4^3
        )
        self.global_pool = nn.AdaptiveAvgPool3d(1)
        self.fc = nn.Linear(64, feature_dim)

    def forward(self, x):
        # x: (B, 1, D, H, W)
        x = self.features(x)       # (B, 64, 4, 4, 4)
        x = self.global_pool(x)    # (B, 64, 1, 1, 1)
        x = x.view(x.size(0), -1)  # (B, 64)
        return self.fc(x)          # (B, feature_dim)

# 最小调用入口
if __name__ == "__main__":
    # 生成随机点云(模拟单物体表面采样)
    np.random.seed(42)
    num_points = 2048
    # 在球壳上采样,确保体素非空
    phi = np.random.uniform(0, 2*np.pi, num_points)
    theta = np.random.uniform(0, np.pi, num_points)
    r = 0.5 + 0.1 * np.random.randn(num_points)
    x = r * np.sin(theta) * np.cos(phi)
    y = r * np.sin(theta) * np.sin(phi)
    z = r * np.cos(theta)
    points = np.stack([x, y, z], axis=1)

    # 体素化
    voxel = point_cloud_to_voxel(points, grid_size=32, max_range=1.0)
    print(f"体素网格形状: {voxel.shape}")          # (1, 1, 32, 32, 32)
    print(f"非空体素数量: {voxel.sum().item():.0f}")

    # 特征提取
    model = VoxelFeatureExtractor(input_channels=1, feature_dim=128)
    model.eval()
    with torch.no_grad():
        feature = model(voxel)
    print(f"输出特征形状: {feature.shape}")          # (1, 128)
    print(f"特征L2范数: {feature.norm().item():.4f}")

运行与验证

  1. 依赖:Python 3.8+,PyTorch 1.9+,NumPy。无需额外数据集。
  2. 运行方式:将上述代码保存为 voxel_feature_extractor.py,执行 python voxel_feature_extractor.py。
  3. 通过条件:
    • 程序无异常退出,输出体素网格形状为 (1, 1, 32, 32, 32)。
    • 非空体素数量大于 0(通常约 200~800,取决于随机采样)。
    • 输出特征形状为 (1, 128),且特征范数为有限非零值。

关键说明

  • 体素化采用二值占用表示(0/1),坐标裁剪至 [-1, 1] 后线性映射到体素索引。对于超出范围的离群点直接裁剪,保证网格索引合法。
  • 3D CNN 采用三层卷积+池化,将 32³ 输入逐步降采样至 4³,再经全局平均池化得到 128 维特征向量。该特征可直接拼接坐标编码后输入 Occupancy 网络。
  • 若需处理真实点云(如 ShapeNet),需先将点云归一化到 [-1, 1] 立方体内,再调用 point_cloud_to_voxel。