3D点云数据缺失补全实战：基于深度学习的点云修复算法与代码实现

1次阅读

没有评论

共计 3027 个字符，预计需要花费 8 分钟才能阅读完成。

在 3D 视觉任务中，点云数据标注的质量直接影响模型的最终性能。然而，由于传感器噪声、物体遮挡或采集设备限制，点云数据经常出现局部缺失。这种缺失会导致：

目标检测模型漏检或定位偏差
分割网络产生破碎的预测结果
配准算法收敛到错误极值点

传统解决方案（如泊松重建或径向基函数插值）往往只能恢复表面连续区域，对复杂结构和大面积缺失束手无策。

最近邻插值：导致点云边缘钝化
MLS 平滑：过度拟合噪声点
基于模板的方法：需要预设物体类别

PointNet++
优点：层级特征提取，支持非均匀点云
缺点：全局特征融合能力弱
PCN (Point Completion Network)
优点：粗到细的双阶段预测，参数量小
缺点：细节恢复能力有限
GRNet
优点：3D 卷积处理体素化特征，保留拓扑结构
缺点：计算成本高

综合计算效率和实现难度，我们选择 PCN 作为基础架构。

import torch
from torch_geometric.data import Data

def preprocess_pcd(partial_pcd, missing_ratio=0.3):
    """
    模拟点云缺失并归一化
    Args:
        partial_pcd: (N,3) tensor
        missing_ratio: 随机删除的比例
    Returns:
        Data(x=features, pos=coordinates)
    """
    N = partial_pcd.shape[0]
    keep_idx = torch.randperm(N)[:int(N*(1-missing_ratio))]
    partial_pcd = partial_pcd[keep_idx]

    # 归一化到单位球
    center = partial_pcd.mean(dim=0)
    partial_pcd -= center
    scale = partial_pcd.abs().max()
    partial_pcd /= scale

    return Data(x=torch.ones(partial_pcd.shape[0],1), 
                pos=partial_pcd)

import torch.nn as nn
import torch.nn.functional as F

class PCNEncoder(nn.Module):
    """编码器：提取全局特征"""
    def __init__(self):
        super().__init__()
        self.mlp1 = nn.Sequential(nn.Conv1d(3, 128, 1),
            nn.BatchNorm1d(128),
            nn.ReLU())
        self.mlp2 = nn.Sequential(nn.Conv1d(128, 256, 1),
            nn.BatchNorm1d(256),
            nn.ReLU())
        self.mlp3 = nn.Sequential(nn.Conv1d(256, 512, 1),
            nn.BatchNorm1d(512),
            nn.ReLU())

    def forward(self, x):
        # x: (B,N,3)
        x = x.transpose(2,1)  # (B,3,N)
        x = self.mlp1(x)
        x = self.mlp2(x)
        x = self.mlp3(x)
        return torch.max(x, dim=2)[0]  # global feature

class PCNDecoder(nn.Module):
    """解码器：生成完整点云"""
    def __init__(self, num_coarse=1024, num_dense=2048):
        super().__init__()
        self.fc1 = nn.Linear(512, 1024)
        self.fc2 = nn.Linear(1024, num_coarse*3)

        # 细节生成网络
        self.mlp = nn.Sequential(nn.Conv1d(3+2, 512, 1),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Conv1d(512, 512, 1),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Conv1d(512, 3, 1)
        )

    def forward(self, x):
        # 粗粒度生成
        coarse = F.relu(self.fc1(x))
        coarse = self.fc2(coarse).view(-1, 1024, 3)  # (B,1024,3)

        # 细粒度生成
        B = coarse.shape[0]
        grid = torch.meshgrid(torch.linspace(-0.05, 0.05, 2),
            torch.linspace(-0.05, 0.05, 2)
        )
        grid = torch.stack(grid, dim=-1).view(4,2).to(x.device)
        grid = grid.unsqueeze(0).repeat(B,1024,1,1)  # (B,1024,4,2)

        points = coarse.unsqueeze(2).repeat(1,1,4,1)  # (B,1024,4,3)
        condition = torch.cat([points, grid], dim=-1)  # (B,1024,4,5)

        condition = condition.view(B*1024*4, -1).transpose(1,0).unsqueeze(0)
        fine = self.mlp(condition).squeeze().view(B, 4096, 3)

        return coarse, torch.cat([coarse, fine], dim=1)

def chamfer_loss(pred, gt):
    """
    计算倒角距离(Chamfer Distance)
    Args:
        pred: (B,N,3)
        gt: (B,M,3)
    """
    dist = torch.cdist(pred, gt)  # (B,N,M)
    min_dist_12 = torch.min(dist, dim=2)[0].mean()
    min_dist_21 = torch.min(dist, dim=1)[0].mean()
    return min_dist_12 + min_dist_21