StraightPCF：基于 Jittor 的点云去噪

本仓库是计图点云去噪赛题的可复现提交代码。主方法将 CVPR 2024 StraightPCF 迁移到 Jittor，并加入两阶段非等强度推理、重叠 patch 软融合和保守残差校准。仓库也包含本项目完成的 ASDN 与 DSNet Jittor 实现，主结果仍以 StraightPCF 为准。

环境安装

推荐 Linux 或 WSL、Python 3.9 和支持当前 Jittor 版本的 CUDA 环境。

conda create -n jittor python=3.9 -y
conda activate jittor
conda install -c conda-forge gcc=10 gxx=10 libgomp -y
python -m pip install -r requirements.txt

安装后可以检查 Jittor 和 CUDA：

python -c "import jittor as jt; print(jt.__version__, jt.has_cuda)"

CPU 自检可在 StraightPCF 命令末尾加入 --use-cuda 0。正式训练和推理建议使用 GPU；ASDN 和 DSNet 是否使用 GPU 由 task YAML 中的 runtime.use_cuda 控制。

数据准备

数据集不随仓库提交。解压后保持如下结构：

data/
  dataset_train/shapenet/<类别>/<样本>/models/model_normalized.obj
  dataset_test_noisy/shapenet/<类别>/<样本>/noisy.npy
  test_gt/shapenet/<类别>/<样本>/clean.npy       # 仅本地评测需要
datalist/
  train.txt
  validate.txt
  test.txt

数据根目录由 configs/data/*.yaml 的 input_dataset_dir 配置；换位置时只需修改该字段。原始数据、预采样点云和预测结果均不应提交。

预训练权重

权重不纳入 Git 仓库。请将项目训练得到或另行发布的权重放入：

checkpoints/base_checkpoint_17.pkl
checkpoints/cvm_checkpoint_56.pkl
checkpoints/clean2pred_checkpoint_18.pkl

其中前两个用于微调初始化，最后一个用于复现报告中的最终推理。配置中的路径均为仓库根目录的相对路径。

训练

从仓库根目录运行（命令行 --seed 会统一设置 Jittor、NumPy 和 Python 随机种子）：

python run.py \
  --task configs/task/train_straightpcf_clean2pred_ft.yaml \
  --seed 123

训练参数来自 task/model/system/data/transform 五类 YAML；学习率、batch size、epoch、数据路径和模型结构均可在其中修改。checkpoint 默认写入 outputs/。

推理

将最终权重放到上述路径后运行：

python run.py \
  --task configs/task/predict_straightpcf_c2p_ckpt18_seed12_softall.yaml \
  --seed 123

评测

python evaluate.py \
  --pred_dir outputs/submission/final \
  --gt_dir data/test_gt \
  --noisy_dir data/dataset_test_noisy \
  --mesh_dir data/dataset_train

仓库结构

configs/       StraightPCF 训练与推理配置
src/           StraightPCF 数据、模型和训练系统
tools/         预采样、融合、代理集评分等工具
asdn/          ASDN 的 Jittor 实现
dsnet/         DSNet 的 Jittor 实现
data/          数据
datalist/      训练、验证和测试划分
docs/          项目报告源文件
outputs/       日志、权重和结果

结果说明

报告中的最终模型为 clean-to-pred checkpoint 18，使用 [0.9, 0.25] 两阶段步长、 [2, 1] 次迭代、1000 点 patch、12 邻域 seed soft fusion 及 0.92 输出校准。项目最终取得比赛第 20 名。ASDN 与 DSNet 为完整探索实现，但没有用于最终提交。

注：由于并非最终实现，因此仓库中的 ASDN 与 DSNet 仅作展示，其代码经过一系列修改后并未验证可行性。

引用与许可

第三方方法和代码来源见 NOTICE。本仓库代码按 MIT License 发布；数据集和预训练权重遵循其各自来源的许可。