在服务器上用自己的二维/三维数据集(.png/.nrrd等)跑nnUNetv2

前言

这篇是我自己的操作笔记，网上可以搜到好多相似文章，大家综合来看，寻找适合自己的，在配置环境过程中出现的问题可参考这篇文章【新手小白】在Linux服务器或本地IDE，跑深度学习代码指南

目录

前言

1. 下载代码并配置环境

2. 准备数据集

二维数据集（如.png）

三维数据集（如.nrrd）

3. 开始训练

4. 验证和推理

5. 其他

---

1. 下载代码并配置环境

打开vscode或pycharm等编辑器，进入到你想存放代码的文件夹下。

选择1：在终端输入命令，直接下载代码

git clone https://github.com/MIC-DKFZ/nnUNet.git

选择2：在github下载代码压缩包，再手动上传到服务器。

此时目录长这样

 先创建一个新的conda环境并启动环境

conda create -n nnunet python=3.10
conda activate nnunet

再给环境安装pytorch包，先输入如下命令，查看CUDA Version是多少

nvidia-smi

我的是CUDA Version: 12.4 ，故选择低于CUDA12.4的pytorch，如CUDA11.8，终端输入如下命令

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

在终端，进入nnUNet目录后，逐步输入以下命令（也可以手动下载后上传服务器：hiddenlayer）

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -e .
git clone https://github.com/FabianIsensee/hiddenlayer.git
cd hiddenlayer
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -e . 

---

2. 准备数据集

此处分为两种情况

• 二维数据集（如 .png、jpg 等）
• 三维数据集（如 .nrrd、niii.gz 等）

二维数据集（如.png）

在nnUNet目录下，新建文件夹nnUNetFrame，DATASET，nnUNet_raw；nnUNet_preprocessed；nnUNet_results，nnUnet_raw_data，以及test；train；valid，在这三个目录下分别创建 images；labels

此时目录长这样

接下来上传 训练、测试、验证数据集到 images 和 labels，我用的都是.png格式的图片，images里放原始图像（shape：512, 512, 3），labels放mask掩码图像（shape: 512, 512）

然后，终端输入如下内容，都是咱刚刚创建好的目录（这一步骤，每重新打开一次电脑都要再运行一遍）

export nnUNet_raw="换成你自己的路径/nnUNet/nnUNetFrame/DATASET/nnUNet_raw"
export nnUNet_preprocessed="换成你自己的路径/nnUNet/nnUNetFrame/DATASET/nnUNet_preprocessed"
export nnUNet_results="换成你自己的路径/nnUNet/nnUNetFrame/DATASET/nnUNet_results"

接下来，到这个路径下：nnUNet/nnunetv2/dataset_conversion找到Dataset120_RoadSegmentation.py

可以把如下内容全部复制替换掉原内容，再按照代码里面的注释，改成自己的内容，并运行

import multiprocessing
import shutil
from multiprocessing import Pool

from batchgenerators.utilities.file_and_folder_operations import *

from nnunetv2.dataset_conversion.generate_dataset_json import generate_dataset_json
from nnunetv2.paths import nnUNet_raw
from skimage import io
from acvl_utils.morphology.morphology_helper import generic_filter_components
from scipy.ndimage import binary_fill_holes