前提概要:
淘论文发现nnDetection框架对肺结节的检测效果挺好,便跑了跑复现了下, 作者诚不欺人,确实挺好, 并附上我依据文档训练的模型。
复现步骤:
1.安装 cuda11.4, cudnn8.2.4
2.创建虚拟环境
conda create -n py38 python=3.8
3.安装 pytorch
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu114
4.下载 nnDetection 项目
进入 nnDetection-main, 执行如下命令, 安装依赖包
pip install -e .
5.设置环境变量
export det_data=/data/code/nnDetection-main/
export det_models=/data/code/nnDetection-main/Task016_Luna/models/
export OMP_NUM_THREADS=1
export det_num_threads=6
6.数据集构造
创建 Task016_Luna/raw 文件夹
将 subset1 - subset10 luna16 数据放到 Task016_Luna/raw 文件夹中
最后文件视图如下格式:
7.预处理
# 该脚本将 luna16 原 .raw 格式数据另存为 .nii.gz 格式, 并利用 annotations.csv 中提供的结节信息,构造 结节的标签信息以及mask数据
python Task016_Luna/scripts/prepare.py
8.依次执行如下官网提供的如下指令
nndet_prep Task016_Luna -np 16 -npp 16
nndet_unpack ${det_data}/Task016_Luna/preprocessed/D3V001_3d/imagesTr 6 (耗时一个半小时)
nndet_train Task016_Luna --sweep
nndet_sweep Task016_Luna RetinaUNetV001_D3V001_3d 0
nndet_eval Task016_Luna RetinaUNetV001_D3V001_3d 0 --boxes --analyze_boxes
nndet_consolidate Task016_Luna RetinaUNetV001_D3V001_3d --sweep_boxes
nndet_predict Task016_Luna RetinaUNetV001_D3V001_3d --fold 0 --test_split --no_preprocess
跑完如上脚本后, 在配置的det_models
路径下会生成如上文件夹, 因为训练实在是太耗时, 我使用RTX5000, 16GB显卡训练了6天得出了10折交叉验证中的一折的结果, 结果已经很不错了, 就没继续跑,如下是框架生成的FROC曲线,可以看到结果确实很不错。
附件(点击下载)
附件中的文件在直接使用模型推理时需要用到