赛道二点云降噪 —— 基于 StraightPCF 的三阶段耦合速度场降噪

赛道：赛道二：基于深度学习的三维点云降噪任务
团队名：铲屎官

本仓库为「计图（Jittor）点云降噪赛题」的参赛代码。任务是给定从三维物体表面采样并加入噪声的点云，预测每个点的位移向量，将含噪点「推回」真实物体表面附近，输出与输入点数严格一致的降噪点云。最终成绩由 Chamfer Distance (CD) 与 Point-to-Surface (P2S) 两个指标各占 50% 综合得到。

快速开始（Quick Start）

完成「环境安装」与「数据准备」后，两条命令即可完成训练与生成提交文件：

# 训练（三阶段，共 300 epoch，权重保存在 experiments/cvm/）
python run.py --task configs/task/train_cvm.yaml

# 推理（加载 experiments/cvm/best_model.pkl，结果写入 my_submission/）
python run.py --task configs/task/predict_cvm.yaml

各步骤的环境、数据、打包提交与评测细节见下文。

方法概述

本方案在赛题提供的 Jittor Baseline 基础上，复现并改进了 CVPR 2024 论文 StraightPCF: Straight Point Cloud Filtering。核心思路是将降噪建模为沿「直线路径」把含噪点搬运回干净表面的最优传输过程，并采用三阶段训练：

Phase 1：训练单个速度模块 vm0，回归常速度位移 δ = X₁ − X₀（干净 − 含噪）。
Phase 2：由 vm0 初始化 vm1，两个速度模块耦合联合训练（损失 L_B），使滤波轨迹更直。
Phase 3：冻结两个速度模块，仅训练 DistanceModule（损失 L_C），预测一个距离标量 d 来缩放步长，避免过冲。

推理时点云被切分为重叠的 1000 点 patch，由耦合速度模块在 DistanceModule 标量缩放下做 Euler 迭代降噪，再拼接还原。

代码主要结构：

src/model/vm.py：速度模块、耦合推理、基于 patch 的降噪。
src/model/feature.py：DGCNN 特征提取器与解码器。
src/system/vm.py：三阶段训练控制器（CVMSystem）。
src/data/augment.py：patch 采样、加噪与插值时间元数据。
configs/：训练 / 推理的配置文件。
run.py：统一入口（设置随机种子并按 --task 分发）。

环境安装

Python 版本：3.9

# 安装计图
conda create -n jittor python=3.9 -y
conda activate jittor
conda install -c conda-forge gcc=10 gxx=10 -y   # 确保 gcc、g++ 版本不高于 10
conda install -c conda-forge libgomp -y         # 确保 OpenMP runtime 存在

# 安装依赖
python -m pip install -r requirements.txt
# 本地评测 P2S 需额外安装 point-cloud-utils（评测可选）
pip install point-cloud-utils

依赖列表见 requirements.txt：jittor、numpy、trimesh、scipy、omegaconf。

数据准备

数据集不随仓库提交，请自行从赛题页面下载并解压到本目录下。

训练集（干净网格）：

tar xzf dataset_train.tar.gz

解压后目录结构：

dataset_train/shapenet/<synset_id>/<model_id>/models/model_normalized.obj

测试集（含噪点云）：

unzip dataset_test_noisy.zip

解压后目录结构：

dataset_test_noisy/shapenet/<synset_id>/<model_id>/noisy.npy

数据根目录通过配置文件中的 input_dataset_dir 字段指定（训练为 dataset_train，推理为 dataset_test_noisy，见 configs/data/train.yaml 与 configs/data/predict.yaml）。使用的样本列表位于 datalist/（train.txt / validate.txt / test.txt）。

训练

完整三阶段训练（共 300 epoch，Phase 1 / 2 / 3 各 100 epoch）：

python run.py --task configs/task/train_cvm.yaml

训练权重保存在 experiments/cvm/ 下（由 configs/system/cvm.yaml 中的 ckpt_save_dir 指定）。
每个阶段会保存该阶段验证集上最优的权重，并将整体最优权重写入 experiments/cvm/best_model.pkl。
关键超参（见 configs/model/cvm.yaml 与 configs/task/train_cvm.yaml）：frame_knn=16、feat_embedding_dim=256、decoder_hidden_dim=64、num_train_points=128、K=2、N=3、λ₁=10.0、λ₂=200.0，Adam，固定学习率 1e-4（不使用学习率调度器）。

推理（生成提交文件）

确认 configs/task/predict_cvm.yaml 中的 load_ckpt 指向你的最佳权重（默认 experiments/cvm/best_model.pkl）。
运行推理：
```
python run.py --task configs/task/predict_cvm.yaml
```
降噪结果保存在 my_submission/（由配置中的 save_dir 指定），目录结构与测试集一致：
```
my_submission/shapenet/<synset_id>/<model_id>/denoised.npy   # float32, shape (N, 3)
```
推理强度当前由代码中的常量 num_repeats=2 控制（硬编码于 src/model/vm.py）；输出点数与输入严格一致（缺失点回退保留其原坐标，见 patch_based_denoise）。

打包为提交文件：

cd my_submission
zip -r ../result.zip shapenet/
cd ..

本地评测

赛题的干净点云与原始网格由组委会持有，本地完整评测需要 GT 数据。若已有 GT，可使用组委会提供的 evaluate.py：

python evaluate.py \
    --pred_dir ./my_submission \
    --gt_dir ./test_gt \
    --noisy_dir ./dataset_test_noisy \
    --mesh_dir ./dataset_train \
    --workers 8

省略 --mesh_dir 时仅计算 CD 得分（不计算 P2S）。

结果说明

指标定义：以含噪输入为 0 分基线，单样本得分为 score = clamp(100 × (1 − dist_pred / dist_noisy), 0, 100)，分别对 CD 与 P2S 计算，最终成绩为全体测试样本上 0.5 × CD_score + 0.5 × P2S_score 的全局平均。

本方案最佳提交结果：

指标	数值
总分	72.63
CD Score	60.31
P2S Score	84.96

对应的平均距离：mean_CD_pred ≈ 0.98e-4（noisy ≈ 2.46e-4），mean_P2S_pred ≈ 0.61e-4（noisy ≈ 1.96e-4）。

复现说明：由于官方分数是对逐样本「截断归一化改进量」取平均，直接用上述全局平均距离取比值并不能精确还原 CD/P2S 得分，二者仅作诊断参考。受 GPU 算子非确定性、硬件/驱动差异以及随机种子影响，复现结果可能与最终成绩存在合理差异。run.py 在启动时统一设置随机种子（--seed，默认 123，同时作用于 Jittor、NumPy、Python random），可通过 python run.py --task ... --seed <n> 指定。

第三方引用与许可

本方案在赛题官方 Jittor Baseline 基础上开发，并参考、复现了 StraightPCF（CVPR 2024）的方法。第三方来源与许可说明见 NOTICE；本仓库许可见 LICENSE。

Edirimuni, D. d. S., Lu, X., Li, G., Wei, L., Robles-Kelly, A., & Li, H. (2024). StraightPCF: Straight Point Cloud Filtering. CVPR 2024, pp. 20721–20730.