【基于nnunet的深度学习医学图像分割网络的使用经验系列】_其一_ResEnc nnU-Net

原创

已于 2025-03-04 20:37:59 修改 · 1k 阅读

14 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #网络 #人工智能 #图像处理 #python

于 2025-02-17 19:01:39 首次发布

文章目录

前言

nnU-Net已经成为了很多医学图像分割网络的baseline。在nnU-Net的开源精神的指引下，一大批成熟的医学图像分割网络也纷纷开源，这对于医学图像分割领域的科研工作者是极大地利好。考虑到国内还没有成熟的nnU-Net的使用踩坑记录，本系列博客应运而生。以此记录我在使用nnU-Net和基于nnU-Net框架的各种开源医学图像分割网络的使用记录。

准备工作

在具备能跑深度学习的NVIDIA显卡的计算机上，安装Ubuntu系统，安装NVIDIA驱动，安装cuda，安装cudnn，安装anaconda，安装vscode，完成基础的环境配置。（这部分内容请自行搜索）

也可以考虑租用云服务器，具体有关于云服务器炼丹的方法请自行搜索。此时，就不必创建虚拟环境了，在配置好cuda的主机上启动后，直接进行准备工作的第三步，安装nnU-Net v2即可

本文讲述nnU-Netv2的使用，基于nnU-Net框架的各种开源医学图像分割网络的使用请期待后文。

创建虚拟环境

conda create -n new_env python=3.10

激活虚拟环境

conda activate new_env

安装nnU-Net v2

pip install nnunetv2

安装nnunetv2时，pytorch等相关包会自动安装，耐心等待即可。

安装好后，创建三个文件夹，分别为

nnUNet_raw
nnUNet_preprocessed
nnUNet_results

其中nnUNet_raw用于存放原始医学影像数据，nnUNet_preprocessed用于存放自动生成的网络结构和训练相关参数，nnUNet-results用于存放训练和测试结果。

找到nnunetv2的安装目录，一般为

/home/***/miniconda3/envs/new_env/lib/python3.10/site-packages/nnunetv2

如果直接装在root中，使用which pip查找即可，此时可能存在位置例如：

/root/miniconda3/lib/python3.10/site-packages/nnunetv2

其中***为ubuntu系统的账户名
同时，上述所有new_env都可以改成任意anaconda的虚拟环境名称。

在该文件夹中，找到paths.py文件，打开，将代码

nnUNet_raw=
nnUNet_preprocessed=
nnUNet_results=

的右边，修改为刚刚创建的三个文件夹的路径。

至此前期准备工作完成。

制作训练集

基于nnU-Net框架的各种开源医学图像分割网络可以针对使用的医学图像分割数据库，进行全自动的deep-supervised的patch-based的网络设计，网络训练，以及验证集测试。

如果我们仅想在自己的数据集的基础上跑一跑别人的代码，不修改网络结构或者调整训练参数。那么，就可以直接将验证集视为测试集。这样一遍跑下来，网络训练和测试都完成了。

因此，使用nnU-Net框架的各种开源医学图像分割网络的第一要务，就是制作符合规范的数据集。

以下通过nnU-Net v2的官网给定的制作Dataset042_BraTS18的代码，讲解如何制作基于nnU-Net框架的各种开源医学图像分割网络的数据集。
代码的链接如下：
https://github.com/MIC-DKFZ/nnUNet/blob/master/nnunetv2/dataset_conversion/Dataset042_BraTS18.py

代码如下：

import multiprocessing
import shutil
import SimpleITK as sitk
import numpy as np
from tqdm import tqdm
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


def copy_BraTS_segmentation_and_convert_labels_to_nnUNet(in_file: str, out_file: str) -> None:
    # use this for segmentation only!!!
    # nnUNet wants the labels to be continuous. BraTS is 0, 1, 2, 4 -> we make that into 0, 1, 2, 3
    img = sitk.ReadImage(in_file)
    img_npy = sitk.GetArrayFromImage(img)

    uniques = np.unique(img_npy)
    for u in uniques:
        if u not in [0, 1, 2, 4]:
            raise RuntimeError('unexpected label')

    seg_new = np.zeros_like(img_npy)
    seg_new[img_npy == 4] = 3
    seg_new[img_npy == 2] = 1
    seg_new[img_npy == 1] = 2
    img_corr = sitk.GetImageFromArray(seg_new)
    img_corr.CopyInformation(img)
    sitk.WriteImage(img_corr, out_file)


def convert_labels_back_to_BraTS(seg: np.ndarray):
    new_seg = np.zeros_like(seg)
    new_seg[seg == 1] = 2
    new_seg[seg == 3] = 4
    new_seg[seg == 2] = 1
    return new_seg


def load_convert_labels_back_to_BraTS(filename, input_folder, output_folder):
    a = sitk.ReadImage(join(input_folder, filename))
    b = sitk.GetArrayFromImage(a)
    c = convert_labels_back_to_BraTS(b)
    d = sitk.GetImageFromArray(c)
    d.CopyInformation(a)
    sitk.WriteImage(d, join(output_folder, filename))


def convert_folder_with_preds_back_to_BraTS_labeling_convention(input_folder: str, output_folder: str,
                                                                num_processes: int = 12):
    """
    reads all prediction files (nifti) in the input folder, converts the labels back to BraTS convention and saves the
    """
    maybe_mkdir_p(output_folder)
    nii = subfiles(input_folder, suffix='.nii.gz', join=False)
    with multiprocessing.get_context("spawn").Pool(num_processes) as p:
        p.starmap(load_convert_labels_back_to_BraTS, zip(nii, [input_folder] * len(nii), [output_folder] * len(nii)))


if __name__ == '__main__':
    brats_data_dir = ...

    task_id = 42
    task_name = "BraTS2018"

    foldername = "Dataset%03.0d_%s" % (task_id, task_name)

    # setting up nnU-Net folders
    out_base = join(nnUNet_raw, foldername)
    imagestr = join(out_base, "imagesTr")
    labelstr = join(out_base, "labelsTr")
    maybe_mkdir_p(imagestr)
    maybe_mkdir_p(labelstr)

    case_ids_hgg = subdirs(join(brats_data_dir, "HGG"), prefix='Brats', join=False)
    case_ids_lgg = subdirs(join(brats_data_dir, "LGG"), prefix="Brats", join=False)

    print("copying hggs")
    for c in tqdm(case_ids_hgg):
        shutil.copy(join(brats_data_dir, "HGG", c, c + "_t1.nii"), join(imagestr, c + '_0000.nii'))
        shutil.copy(join(brats_data_dir, "HGG", c, c + "_t1ce.nii"), join(imagestr, c + '_0001.nii'))
        shutil.copy(join(brats_data_dir, "HGG", c, c + "_t2.nii"), join(imagestr, c + '_0002.nii'))
        shutil.copy(join(brats_data_dir, "HGG", c, c + "_flair.nii"), join(imagestr, c + '_0003.nii'))

        copy_BraTS_segmentation_and_convert_labels_to_nnUNet(join(brats_data_dir, "HGG", c, c + "_seg.nii"),
                                                             join(labelstr, c + '.nii'))
    print("copying lggs")
    for c in tqdm(case_ids_lgg):
        shutil.copy(join(brats_data_dir, "LGG", c, c + "_t1.nii"), join(imagestr, c + '_0000.n