nnUnet使用2d數據訓練方法-DKFZ官方版

nnUnet使用2d數據訓練方法-DKFZ官方版

上一篇文章介紹了《保姆級教程：nnUnet在2維圖像的訓練和測試》，採用的是自己的2d數據集進行2d到3d的數據轉換，內容包括nnUnet介紹、環境配置、數據配置、預處理、訓練過程、確定最佳的U-Net配置、運行推斷，算是帶著大家在2d數據情況下把nnUnet訓練和測試了一遍。

最近官方也更新了nnUnet在2d數據情況下的訓練方法，連結為：https://github.com/MIC-DKFZ/nnUNet/blob/master/documentation/dataset_conversion.md，將2d的麻薩諸塞州道路數據road_segmentation_ideal轉換成3d的Task120_MassRoadsSeg數據。

麻薩諸塞州道路數據集是衛星標註圖像，從航空圖像中分割道路是一項具有挑戰性的任務。來自附近樹木的障礙物、相鄰建築物的陰影、道路紋理和顏色的變化、道路等級的不平衡（由於道路圖像像素相對較少）是阻礙當前模型分割從圖像一端延伸到另一端的尖銳道路邊界的其他挑戰。高質量的航空圖像數據集有助於對現有方法進行比較，並在機器學習和計算機視覺領域引起人們對航空圖像應用的興趣。

輸入圖像如下：

標註圖像如下：

一、官方說明How to use 2D data with nnU-Net

nnU-Net was originally built for 3D images. It is also strongestwhen applied to 3D segmentation problems because a large proportion of itsdesign choices were built with 3D in mind. Also note that many 2D segmentationproblems, especially in the non-biomedical domain, may benefit from pretrainednetwork architectures which nnU-Net does not support. Still, there is certainlya need for an out of the box segmentation solution for 2D segmentationproblems. And also on 2D segmentation tasks nnU-Net cam perform extremely well!We have, for example, won a 2D task in the cell tracking challenge with nnU-Net(see our Nature Methods paper) and we have also successfully applied nnU-Net tohistopathological segmentation problems. Working with 2D data in nnU-Netrequires a small workaround in the creation of the dataset. Essentially, allimages must be converted to pseudo 3D images (so an image with shape (X, Y)needs to be converted to an image with shape (1, X, Y). The resulting imagemust be saved in nifti format. Hereby it is important to set the spacing of thefirst axis (the one with shape 1) to a value larger than the others. If you areworking with niftis anyways, then doing this should be easy for you. Thisexample here is intended for demonstrating how nnU-Net can be used with'regular' 2D images. We selected the massachusetts road segmentation datasetfor this because it can be obtained easily, it comes with a good amount oftraining cases but is still not too large to be difficult to handle.

See here for anexample. This script contains a lot of comments and useful information. Alsohave a look here.

How to convert other image formats tonifti

Please have a look at the following tasks:

二、Task120數據集

這裡Task120的連結為：

https://github.com/MIC-DKFZ/nnUNet/blob/master/nnunet/dataset_conversion/Task120_Massachusetts_RoadSegm.py

首先簡單介紹一下代碼的功能：

創建原始數據的文件夾Task120_MassRoadsSeg以及子文件夾imagesTr、imagesTs、labelsTs、labelsTr。

/road_segmentation_ideal路徑內的有訓練集training和測試集testing兩個文件夾，數據集文件夾內有輸入圖像input和標籤output兩個文件夾，分別提取到每個數據的名稱。

調用convert_2d_image_to_nifti函數，讀取每個數據，將3通道的圖像拆成3個模態的3個數據，設置每個像素點實際的長度與寬度Spacing，並保存成nii.gz的3維數據，這裡的3維數據其實只有1層。

原始圖像img-1.png如下：

生成的3個通道的3個模態數據img-1_0000.nii.gz、img-1_0001.nii.gz和img-1_0002.nii.gz如下，是用ITK-SNAP打開的，只展示沿z軸的一層圖像。

三、代碼詳解

此部分代碼文件為：

nnUNet/nnunet/dataset_conversion/Task120_Massachusetts_RoadSegm.py

1. 首先導入需要的包

import numpy as np
from batchgenerators.utilities.file_and_folder_operations import *
from nnunet.dataset_conversion.utils import generate_dataset_json
from nnunet.paths import nnUNet_raw_data, preprocessing_output_dir
from nnunet.utilities.file_conversions impo

2. 定義main函數並做介紹

if __name__ == '__main__':
    """
    nnU-Net was originally built for 3D images. It is also strongest when applied to 3D segmentation problems because a 
    large proportion of its design choices were built with 3D in mind. Also note that many 2D segmentation problems, 
    especially in the non-biomedical domain, may benefit from pretrained network architectures which nnU-Net does not
    support.
    Still, there is certainly a need for an out of the box segmentation solution for 2D segmentation problems. And 
    also on 2D segmentation tasks nnU-Net cam perform extremely well! We have, for example, won a 2D task in the cell 
    tracking challenge with nnU-Net (see our Nature Methods paper) and we have also successfully applied nnU-Net to 
    histopathological segmentation problems. 
    Working with 2D data in nnU-Net requires a small workaround in the creation of the dataset. Essentially, all images 
    must be converted to pseudo 3D images (so an image with shape (X, Y) needs to be converted to an image with shape 
    (1, X, Y). The resulting image must be saved in nifti format. Hereby it is important to set the spacing of the 
    first axis (the one with shape 1) to a value larger than the others. If you are working with niftis anyways, then 
    doing this should be easy for you. This example here is intended for demonstrating how nnU-Net can be used with 
    'regular' 2D images. We selected the massachusetts road segmentation dataset for this because it can be obtained 
    easily, it comes with a good amount of training cases but is still not too large to be difficult to handle.
"""

大致意思是nnUnet最初適用於3D數據，當處理2D數據時需要將數據做成一個偽3D的圖像，形狀為（X，Y）的圖像需要轉換為形狀為（1，X，Y）的圖像，結果圖像必須以nifti格式保存，將第一軸（具有形狀1的軸）的間距設置為大於其他軸的值。該數據集為2d數據麻薩諸塞州道路分割數據集。

3. 下載數據集、修改文件路徑並創建nnUnet能識別的文件夾

   # download dataset from https://www.kaggle.com/insaff/massachusetts-roads-dataset
    # extract the zip file, then set the following path according to your system:
    base = '/media/fabian/data/road_segmentation_ideal'
    # this folder should have the training and testing subfolders

    # now start the conversion to nnU-Net:
    task_name = 'Task120_MassRoadsSeg'
    target_base = join(nnUNet_raw_data, task_name)
    target_imagesTr = join(target_base, "imagesTr")
    target_imagesTs = join(target_base, "imagesTs")
    target_labelsTs = join(target_base, "labelsTs")
    target_labelsTr = join(target_base, "labelsTr")

    maybe_mkdir_p(target_imagesTr)
    maybe_mkdir_p(target_labelsTs)
    maybe_mkdir_p(target_imagesTs)
maybe_mkdir_p(target_labelsTr)

這裡給出了數據集的下載連結：

https://www.kaggle.com/insaff/massachusetts-roads-dataset，解壓並修改文件路徑base。

創建原始數據的文件夾Task120_MassRoadsSeg以及子文件夾imagesTr、imagesTs、labelsTs、labelsTr。

4.  對訓練數據文件進行遍歷

# convert the training examples. Not all training images have labels, so we just take the cases for which there are
    # labels
    labels_dir_tr = join(base, 'training', 'output')
    images_dir_tr = join(base, 'training', 'input')
    training_cases = subfiles(labels_dir_tr, suffix='.png', join=False)
    for t in training_cases:
        unique_name = t[:-4]  # just the filename with the extension cropped away, so img-2.png becomes img-2 as unique_name
        input_segmentation_file = join(labels_dir_tr, t)
        input_image_file = join(images_dir_tr, t)

        output_image_file = join(target_imagesTr, unique_name)  # do not specify a file ending! This will be done for you
        output_seg_file = join(target_labelsTr, unique_name)  # do not specify a file ending! This will be done for you

這裡提示一下，該數據集部分數據無標籤，所以代碼查找的數據以有label的數據為準，無label的數據直接忽略。

針對/road_segmentation_ideal路徑內的訓練集training文件夾，數據集文件夾內有輸入圖像input和標籤output兩個文件夾，輸入圖像文件夾為images_dir_tr，標籤圖像文件夾為labels_dir_tr，提取到output文件夾裡的所有數據名稱training_cases。

分別依次讀取每個數據（這裡是3通道的彩色2d圖像），獲取數據名稱unique_name（去掉後綴），輸入圖像文件為input_image_file，標註圖像文件為input_segmentation_file，輸出圖像文件為output_image_file（不帶後綴，後面會分成3個模態3個輸出圖像文件），輸出標註文件output_seg_file。

5.  調用convert_2d_image_to_nifti函數進行數據轉換

# this utility will convert 2d images that can be read by skimage.io.imread to nifti. You don't need to do anything.
        # if this throws an error for your images, please just look at the code for this function and adapt it to your needs
        convert_2d_image_to_nifti(input_image_file, output_image_file, is_seg=False)

        # the labels are stored as 0: background, 255: road. We need to convert the 255 to 1 because nnU-Net expects
        # the labels to be consecutive integers. This can be achieved with setting a transform
        convert_2d_image_to_nifti(input_segmentation_file, output_seg_file, is_seg=True,
                                  transform=lambda x: (x == 255).astype(int))

上面的代碼是對訓練數據的輸入圖像和標籤圖像做轉換，調用了convert_2d_image_to_nifti函數，該函數首先做了定義和一些介紹：

from typing import Tuple, List
from skimage import io
import SimpleITK as sitk
import numpy as np
import tifffile


def convert_2d_image_to_nifti(input_filename: str, output_filename_truncated: str, spacing=(999, 1, 1),
                              transform=None, is_seg: bool = False) -> None:
    """
    Reads an image (must be a format that it recognized by skimage.io.imread) and converts it into a series of niftis.
    The image can have an arbitrary number of input channels which will be exported separately (_0000.nii.gz,
    _0001.nii.gz, etc for images and only .nii.gz for seg).
    Spacing can be ignored most of the time.
    !!!2D images are often natural images which do not have a voxel spacing that could be used for resampling. These images
    must be resampled by you prior to converting them to nifti!!!
    Datasets converted with this utility can only be used with the 2d U-Net configuration of nnU-Net
    If Transform is not None it will be applied to the image after loading.
    Segmentations will be converted to np.uint32!
    :param is_seg:
    :param transform:
    :param input_filename:
    :param output_filename_truncated: do not use a file ending for this one! Example: output_name='./converted/image1'. This
    function will add the suffix (_0000) and file ending (.nii.gz) for you.
    :param spacing:
    :return:
    """

大致意思是要做數據轉換了，輸入圖像可以是任意通道的（通常是3通道），後面會單獨輸出_0000.nii.gz、_0001.nii.gz等的結果，但標籤數據還是輸出不帶模態標誌的.nii.gz，間距spacing可以忽略，轉換之後的數據集只能用2d unet配置使用。

首先用skimage讀圖

img = io.imread(input_filename)

如果是標籤圖像，則應用transform，將[0,255]的標籤圖像轉換成[0,1]的標籤圖像

if transform is not None:
    img = transform(img)

這裡的

transform=lambda x: (x == 255).astype(int))

如果是1通道的2d圖像，則圖像形狀可能是（256，256），需要增加2個維度，變成（1，1，256，256）

if len(img.shape) == 2:  # 2d image with no color channels
    img = img[None, None]  # add dimensions

如果是3通道的2d圖像，則圖像形狀可能是（256，256，3），先轉置成（3，256，256），然後增加1個維度，變成(3, 1, 256, 256)。

else:
    assert len(img.shape) == 3, "image should be 3d with color channel last but has shape %s" % str(img.shape)
    # we assume that the color channel is the last dimension. Transpose it to be in first
    img = img.transpose((2, 0, 1))
    # add third dimension
    img = img[:, None]

再次確認標籤圖像的通道數為1

# image is now (c, x, x, z) where x=1 since it's 2d
if is_seg:
    assert img.shape[0] == 1, 'segmentations can only have one color channel, not sure what happened here'

遍歷圖像的通道數（通常為3個），分別提取不同通道的圖像，轉換成itk文件，設置spacing層厚，這裡預定義的spacing=(999, 1, 1)，經過list(spacing)[::-1]後變成[1, 1, 999]，sitk圖像順序是x,y,z三個方向的大小，numpy矩陣的順序是z,y,x三個方向的大小，所以spacing需要轉換成[1, 1, 999]。之後針對輸入圖像和標籤圖像分別存儲為帶模態標誌0000/0001/0002的nifti數據和不帶模態標誌的nifti數據。

for j, i in enumerate(img):

    if is_seg:
        i = i.astype(np.uint32)
    itk_img = sitk.GetImageFromArray(i)
    itk_img.SetSpacing(list(spacing)[::-1])
    if not is_seg:
        sitk.WriteImage(itk_img, output_filename_truncated + "_%04.0d.nii.gz" % j)
    else:
        sitk.WriteImage(itk_img, output_filename_truncated + ".nii.gz")

測試數據方法一樣，

# now do the same for the test set
    labels_dir_ts = join(base, 'testing', 'output')
    images_dir_ts = join(base, 'testing', 'input')
    testing_cases = subfiles(labels_dir_ts, suffix='.png', join=False)
    for ts in testing_cases:
        unique_name = ts[:-4]
        input_segmentation_file = join(labels_dir_ts, ts)
        input_image_file = join(images_dir_ts, ts)

        output_image_file = join(target_imagesTs, unique_name)
        output_seg_file = join(target_labelsTs, unique_name)

        convert_2d_image_to_nifti(input_image_file, output_image_file, is_seg=False)
        convert_2d_image_to_nifti(input_segmentation_file, output_seg_file, is_seg=True,
                                  transform=lambda x: (x == 255).astype(int))

最後生成json文件

# finally we can call the utility for generating a dataset.json
    generate_dataset_json(join(target_base, 'dataset.json'), target_imagesTr, target_imagesTs, ('Red', 'Green', 'Blue'),
                          labels={1: 'street'}, dataset_name=task_name, license='hands off!')

最終轉換之後的數據格式為：

nnUNet_raw_data_base/nnUNet_raw_data/Task120_MassRoadsSeg

├── dataset.json

├── imagesTr

│   ├── img-2_0000.nii.gz

│   ├── img-2_0001.nii.gz

│   ├── img-2_0002.nii.gz

│   ├── ...

├── imagesTs

│   ├── img-1_0000.nii.gz

│   ├── img-1_0001.nii.gz

│   ├── img-1_0002.nii.gz

│   ├── ...

├── labelsTr

│   ├── img-2.nii.gz

│   ├── img-7.nii.gz

│   ├── ...

├── labelsTs

│   ├── img-1.nii.gz

│   ├── img-2.nii.gz

│   ├── ...

這樣數據集就處理好了，之後再按照步驟進行預處理、訓練過程、確定最佳的U-Net配置、運行推斷就好了，可以參照官方https://github.com/MIC-DKFZ/nnUNet，也可以參考我們的前一篇文章《保姆級教程：nnUnet在2維圖像的訓練和測試》。

最後2021年衝鴨！！！

掃它！掃它！掃它！