jittor_mlghenshuai_PCD

基于 Jittor 深度学习框架的 Rectified Flow 点云降噪模型。

• 模型架构：DGCNN + FiLM 时间条件注入的 VelocityNet，学习从含噪点云到干净点云的速度场 $v_{\theta }(x_t,t)v\_\backslash theta(x\_t,\; t)$
• 训练：在 ShapeNet 网格表面采样 patch 对（干净/加噪），使用 Rectified Flow loss 训练
• 推理：对含噪点云做重叠 patch 采样 → 反向 Euler 积分 → 加权融合，输出去噪点云
• 评测：Chamfer Distance (CD) + Point-to-Surface (P2S) 双指标，百分制评分

---

1. 环境安装

Python 版本

• Python 3.9（推荐，与 Jittor 兼容性最佳）

安装步骤

# 创建并激活 conda 环境
conda create -n jittor python=3.9 -y
conda activate jittor

# 安装 gcc/g++（版本不高于 10，Jittor 编译需要）
conda install -c conda-forge gcc=10 gxx=10 -y

# 安装 OpenMP runtime
conda install -c conda-forge libgomp -y

# 安装计图
python -m pip install jittor

# 安装其他依赖
python -m pip install -r requirements.txt

注意：Jittor 是元算子框架，首次运行时会对算子进行即时编译（JIT compile），请确保 gcc/g++ 版本不高于 10，且有稳定的网络连接。

---

2. 数据准备

数据来源

训练和评测使用 ShapeNet 数据集的标准三角网格（.obj 格式）。

目录结构

其中 local_val_A_noisy 和 local_val_A_gt 需自行生成

Data/
├── dataset_train/              # 训练网格
│   └── shapenet/
│       └── 02691156/           # 类别
│           └── <model_id>/
│               └── model_normalized.obj
│
├── local_val_A_noisy/          # 验证/测试含噪点云
│   └── shapenet/
│       └── <category>/
│           └── <model_id>/
│               └── noisy.npy
│
└─── local_val_A_gt/             # 验证/测试干净点云
    └── shapenet/
        └── <category>/
            └── <model_id>/
                └── clean.npy

数据列表文件

datalist/ 目录下提供三个文本文件，每行一个模型相对路径：

文件	样本数	用途
new_train.txt	15,733	训练集
new_validate.txt	100	验证集
test.txt	200	测试集

文件中每行格式示例：shapenet/03642806/2134ad3fc25a6284193a4c984002ed32

数据根目录配置

所有路径通过命令行参数传入，无需配置文件：

参数	说明
--mesh_root	训练网格根目录，与 datalist 中的相对路径拼接得到 .obj 文件
--cache_root	网格采样点缓存目录（自动缓存 50k 点的 .npy 文件，加速后续 epoch）
--in_root	推理时含噪点云输入目录
--out_root	推理时降噪结果输出目录
--gt_dir / --noisy_dir / --mesh_dir	评测时 GT / 含噪 / 网格目录

---

3. 训练

一键训练命令

python train_rf.py \
    --list_file datalist/new_train.txt \
    --mesh_root ./Data/dataset_train \
    --cache_root cache/mesh_pts_50k \
    --save_dir experiments/rf_v1 \
    --save_every 1 \
    --batch_size 16 \
    --epochs 30 \
    --lr 2e-4 \
    --patch_size 1024 \
    --num_workers 6

或直接运行脚本：

bash scripts/train.sh

关键参数说明

参数	默认值	说明
--batch_size	16	每批 patch 数
--epochs	30	训练轮数
--lr	2e-4	初始学习率（线性预热 + cosine 衰减至 0.1×）
--patch_size	1024	每个 patch 的点数
--sigma_min / --sigma_max	0.005 / 0.020	加噪标准差范围（单位球归一化后）
--feat_dims	64 128 256 256	EdgeConv 各层通道数
--knn_k	16	kNN 邻居数
--warmup_iters	500	学习率预热步数
--grad_clip	1.0	梯度裁剪阈值
--save_every	2	每隔 N 个 epoch 保存一次 checkpoint
--resume	-	从指定 checkpoint 恢复训练

Checkpoint 格式

Checkpoint 以 Jittor .pkl 格式保存

---

4. 评测 / 推理

推理（生成降噪结果）

python infer_rf.py \
    --ckpt experiments/rf_v1/checkpoint_29.pkl \
    --in_root ./Data/local_val_A_noisy \
    --out_root ./results_submission \
    --outer_iters 2 \
    --n_steps 4 \
    --patch_size 1024 \
    --seed_k 6

或直接运行脚本：

bash scripts/infer.sh

推理参数说明

参数	默认值	说明
--ckpt	必填	训练好的 checkpoint 路径
--in_root	必填	含噪 .npy 文件目录
--out_root	必填	降噪结果输出目录
--n_steps	4	反向 Euler 积分步数（越大越精细，越慢）
--outer_iters	1	外循环迭代次数（重复 patch 采样/融合）
--patch_size	1024	每个 patch 点数
--seed_k	6	每个点的平均覆盖 patch 数（越大融合越平滑）
--schedule	uniform	时间调度策略：uniform / front_dense / back_dense / cosine / power
--batch_patches	8	每批处理的 patch 数（GPU 显存相关）

评测（计算 CD 和 P2S 分数）

python evaluate.py \
    --pred_dir ./results_submission \
    --gt_dir ./Data/local_val_A_gt \
    --noisy_dir ./Data/local_val_A_noisy \
    --mesh_dir ./Data/dataset_train \
    --workers 8

一键推理+评测

bash scripts/eval.sh \
    --ckpt experiments/rf_v1/checkpoint_29.pkl \
    --in_root ./Data/local_val_A_noisy \
    --gt_dir ./Data/local_val_A_gt \
    --noisy_dir ./Data/local_val_A_noisy \
    --mesh_dir ./Data/dataset_train \
    --gpu 0

eval.sh 会自动从 checkpoint 文件名中提取编号，生成 results_jt_rf_<编号> 格式的输出目录。

---

5. 结果说明

指标定义

Chamfer Distance (CD)

$\text{CD}(A,B)=\frac{1}{\mathrm{\mid }A\mathrm{\mid }}\sum _{a\in A}\underset{b\in B}{}\mathrm{\parallel }a-b{\mathrm{\parallel }}^2+\frac{1}{\mathrm{\mid }B\mathrm{\mid }}\sum _{b\in B}\underset{a\in A}{}\mathrm{\parallel }b-a{\mathrm{\parallel }}^2\backslash text\{CD\}(A,\; B)\; =\; \backslash frac\{1\}\{|A|\}\backslash sum\_\{a\; \backslash in\; A\}\; \backslash min\_\{b\; \backslash in\; B\}\; \backslash |a\; -\; b\backslash |^2\; +\; \backslash frac\{1\}\{|B|\}\backslash sum\_\{b\; \backslash in\; B\}\; \backslash min\_\{a\; \backslash in\; A\}\; \backslash |b\; -\; a\backslash |^2$

• 衡量两个点云之间的双向最近邻平均平方距离
• 使用 scipy cKDTree 进行 $O(N\log M)O(N\; \backslash log\; M)$ 加速
• 计算前以参考点云（GT）归一化到单位球，预测点云施加相同变换

Point-to-Surface (P2S)

$\text{P2S}(P,M)=\frac{1}{\mathrm{\mid }P\mathrm{\mid }}\sum _{p\in P}\underset{q\in \text{surface}(M)}{}\mathrm{\parallel }p-q{\mathrm{\parallel }}^2\backslash text\{P2S\}(P,\; M)\; =\; \backslash frac\{1\}\{|P|\}\backslash sum\_\{p\; \backslash in\; P\}\; \backslash min\_\{q\; \backslash in\; \backslash text\{surface\}(M)\}\; \backslash |p\; -\; q\backslash |^2$

• 衡量预测点云中每个点到原始网格表面的最近距离平方的均值
• 优先使用 point-cloud-utils（基于 BVH 加速的精确表面投影）
• 若 point-cloud-utils 未安装，自动回退到 cKDTree 顶点近似（精度略低）

评分映射

对每个测试样本 $ii$：

${\text{cd\_score}}_i=\text{clamp}(100\times (1-\frac{{\text{CD}}_{\text{pred}}}{{\text{CD}}_{\text{noisy}}}),0,100)\backslash text\{cd\backslash \_score\}\_i\; =\; \backslash text\{clamp\}\backslash left(100\; \backslash times\; \backslash left(1\; -\; \backslash frac\{\backslash text\{CD\}\_\backslash text\{pred\}\}\{\backslash text\{CD\}\_\backslash text\{noisy\}\}\backslash right),\backslash \; 0,\backslash \; 100\backslash right)$

${\text{p2s\_score}}_i=\text{clamp}(100\times (1-\frac{{\text{P2S}}_{\text{pred}}}{{\text{P2S}}_{\text{noisy}}}),0,100)\backslash text\{p2s\backslash \_score\}\_i\; =\; \backslash text\{clamp\}\backslash left(100\; \backslash times\; \backslash left(1\; -\; \backslash frac\{\backslash text\{P2S\}\_\backslash text\{pred\}\}\{\backslash text\{P2S\}\_\backslash text\{noisy\}\}\backslash right),\backslash \; 0,\backslash \; 100\backslash right)$

最终得分（所有样本全局平均，不分类别）：

$\text{Final}=0.5\times \overline{\text{cd\_score}}+0.5\times \overline{\text{p2s\_score}}\backslash text\{Final\}\; =\; 0.5\; \backslash times\; \backslash overline\{\backslash text\{cd\backslash \_score\}\}\; +\; 0.5\; \backslash times\; \backslash overline\{\backslash text\{p2s\backslash \_score\}\}$

• 满分 100 分，分数越高表示降噪效果越好
• 得分 ≤ 0 表示降噪后不比含噪输入更好（截断为 0）
• 缺失预测的样本记 0 分

本地测试得分

在本地采用从验证集生成的测试数据进行测试 结果如下：

=================================================================
  点云降噪评测结果
=================================================================
  评测样本总数:       100
  有效预测数:         100
  缺失预测数:         0
  并行进程数:         8
  评测耗时:           6.6s
-----------------------------------------------------------------
  平均 CD_pred:       0.00012856
  平均 CD_noisy:      0.00035816
  CD 得分:            60.54 / 100.00
  平均 P2S_pred:      0.00044500
  平均 P2S_noisy:     0.00066228
  P2S 得分:           81.68 / 100.00
-----------------------------------------------------------------
  最终得分 (0.5×CD + 0.5×P2S):  71.11 / 100.00
=================================================================

与线上提交 / 最终成绩的差异

• 本脚本计算的是本地验证集结果，线上提交使用隐藏测试集，样本不同，分数会有所差异
• 若 point-cloud-utils 未安装，P2S 计算使用顶点近似而非精确表面投影，P2S 分数可能略高于真实值（顶点近似低估了表面距离），建议线上提交前安装 point-cloud-utils 确保一致性
• 最终提交的分数以线上平台返回结果为准

---

6. 后处理（平面离群点修复）

推理输出的降噪点云可能残留少量平面离群点——即局部邻域呈强平面结构、但该点自身偏离平面较远的点。postprocess_planar_outliers.py 对此类点进行检测并投影修复。

原理

1. 对每个点取 kNN 邻域，用 PCA 拟合局部平面（最小特征值对应法向）
2. 若邻域平面性足够强（normal_var ≤ plane_ratio）且该点的 signed distance 超过阈值，则标记为离群点
3. 将离群点沿法向投影回拟合平面（带 blend 系数和最大移动距离约束）
4. 支持多轮迭代修复 + 小簇修复（--cluster_mode）+ 锚点混合（--anchor_blend，向原始含噪点云回退）

使用方式

python postprocess_planar_outliers.py \
    --input_dir ./results_submission \
    --output_dir ./results_submission_planar_fix \
    --filename denoised.npy \
    --workers 8

常用可选参数：

参数	默认值	说明
--k	36	kNN 邻域大小
--plane_ratio	0.075	判定为平面的 λ_min / Σλ 阈值
--dist_mult	2.7	离群距离阈值（× robust_sigma）
--blend	0.75	投影混合系数（1=完全投影，<1 更保守）
--repair_rounds	2	多轮修复次数
--cluster_mode	-	启用小簇离群点修复
--anchor_blend	0	向 noisy 锚点混合的强度（>0 开启）
--gt_dir / --noisy_dir	-	若提供，修复前后自动计算 CD score 对比

---

项目结构

jittor_mlghenshuai_PCD/
├── train_rf.py             # 训练入口
├── infer_rf.py             # 推理入口（生成降噪结果）
├── evaluate.py                    # 评测入口（CD + P2S 评分）
├── postprocess_planar_outliers.py # 后处理（平面离群点修复）
├── model/
│   ├── rf_model.py    # VelocityNet 网络结构（DGCNN + FiLM）
│   └── rf_flow.py     # Rectified Flow loss / 采样 / 时间调度
├── datalist/
│   ├── new_train.txt         # 训练集路径列表
│   ├── new_validate.txt      # 验证集路径列表
│   └── test.txt              # 测试集路径列表
├── scripts/
│   ├── train.sh              # 训练启动脚本
│   ├── infer.sh              # 推理启动脚本
│   └── eval.sh               # 推理+评测一键脚本
├── experiments/              # 训练输出（checkpoint + 日志）
├── cache/                    # 网格采样点缓存
└── README.md