nnUNet使用教程(使用自己的数据训练自己的模型)

1.制作自己的数据集

nnunet作为unet的一种改进,在3d医学图像上使用有非常好的效果,本篇教程主要说明nnunet在2d图像上的应用,也是相对来说更贴近我们平时使用的方向。本文根据自己的实际操作进行复现,一步一步的训练出自己的模型。

首先我们要有自己的数据,也就是原图,这里我使用的是Eiseg作为分割工具,具体的EISEG使用方法可以自行百度,这里要注意我们在手动分割原图的时候要选择生成coco数据集,生成的coco数据集应该是一个coco.json的格式,(当然也可以使用labelme等其他的分割工具处理数据)。我们拿到了原图和coco格式的json文件使用脚本1 对原图和json文件进行处理得到带有分割标记的图片。

脚本一:

import cv2
from pycocotools.coco import COCO
import os
import shutil
import matplotlib.pyplot as plt
import numpy as np

'''
路径参数
'''
# 原coco数据集的路径
dataDir = "./"
# 用于保存新生成的mask数据的路径
savepath = "coco_mask/"

'''
数据集参数
'''
# coco有80类,这里写要进行二值化的类的名字
# 其他没写的会被当做背景变成黑色
classes_names = ['car']  #自行修改标签

datasets_list = ['coco']


# 生成保存路径
# if the dir is not exists,make it,else delete it
def mkr(path):
    if os.path.exists(path):
        shutil.rmtree(path)
        os.mkdir(path)
    else:
        os.mkdir(path)


# 生成mask图
def mask_generator(coco, width, height, anns_list):
    mask_pic = np.zeros((height, width))
    # 生成mask - 此处生成的是4通道的mask图,如果使用要改成三通道,可以将下面的注释解除,或者在使用图片时搜相关的程序改为三通道
    for single in anns_list:
        # print(single)  # 输出每一个标注信息
        mask_single = coco.annToMask(single)
        mask_pic += mask_single
    # coco.showAnns(anns_list)
    # 转化为255
    for row in range(height):
        for col in range(width):
            if mask_pic[row][col] > 0:
                mask_pic[row][col] = 255
    # 转为三通道
    imgs = np.zeros(shape=(height, width, 3), dtype=np.float32)
    imgs[:, :, 0] = mask_pic[:, :]
    imgs[:, :, 1] = mask_pic[:, :]
    imgs[:, :, 2] = mask_pic[:, :]
    imgs = imgs.astype(np.uint8)
    return imgs


到这里我们就得到我们需要的数据了。

2.处理数据并配置nnunet

首先创建一个项目文件夹,这里我起得名字叫nnUNetFrame_final。我们要在这个文件夹下先创建五个文件夹,DATASET,nnUNet-master  (注:这个文件就是我们在github上下载的源代码),road_segmentation_ideal,(注:这三个文件夹是必备的,后面两个文件夹是结合到后面的脚本我自己定义的,看到后面理解之后可以自己进行修改,后面两个文件夹名称如下:newimage,pic_one_channel)。

到现在为止项目根目录下应该有5个文件夹分别是:DATASET,nnUNet-master,road_segmentation_ideal,newimage,pic_one_channel。

如图:

nnUNet使用教程(使用自己的数据训练自己的模型)_第1张图片

接下来在DATASET中创建nnUNet_preprocessed,nnUNet_raw,nnUNet_trained_models三个文件夹,在其中的nnUNet_raw中创建nnUNet_cropped_data,nnUNet_raw_data文件夹。在nnUNet_raw_data文件夹下创建我们想要训练的项目名称,这里举个例子:比如我这次想要分割车,那么我的文件夹名称应该叫做Task01_Car。要分割心脏就叫做Task02_Heart。这里的01,02也对应了后面要用到的进程id,比如Task01_Car,那么我后面要用到的进程id就为1。

设置bashrc:这一步骤的目的可以让nnunet识别到我们创建的DATASET文件路径。具体方法:

在ubuntu中首先在home目录下按ctrl + h,显示隐藏文件,然后找到.bashrc文件,打开并在最后面加上三行:

export nnUNet_raw_data_base="/home/andre/nnUNetFrame/DATASET/nnUNet_raw"

export nnUNet_preprocessed="/home/andre/nnUNetFrame/DATASET/nnUNet_preprocessed" export RESULTS_FOLDER="/home/andre/nnUNetFrame/DATASET/nnUNet_trained_models"
 

这三行分别是nnUNet_raw,nnUNet_preprocessed,nnUNet_trained_models三个文件的路径,具体路径要自己进行一些修改,这个路径一定要是可以找到这个文件的路径,如果是虚拟机或者容器中使用要确认好地址。

加入三行路径之后点击保存。然后别忘了在home下打开终端,输入source .bashrc更新一下文档。

(处理数据重点!!)继续在根目录下的road_segmentation_ideal文件夹下创建,testing,training两个文件夹,在其中的testing文件夹下创建input和output文件夹,同理在training下也创建input和output文件夹。这里解释一下testing文件顾名思义就是我们的测试数据集,training就是训练数据集,至于他们的比例为多少可以自己定义,其中input为我们的原图,output为单通道的带有分割效果的图片。上一步我们已经获得到了原图和带有标记效果的图片,下面我们使用脚本先把带有标记效果的图片转换成单通道图片,脚本如下:

import os

num = 0
for i in os.listdir("road_segmentation_ideal/training/output"):
  image_path=("road_segmentation_ideal/training/output/"+i)
  #print(image_path)
  num+=1
  img = cv2.imread(image_path,0)
  print(img.shape)
  thresh1, dst = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)
  cv2.imwrite("pic_one_channel/"+i[:-4]+".jpg",dst)
print(num)

看到这里大家就明白之前在根目录下创建的pic_one_channel的作用了。

运行过后我们在pic_one_channel中已经得到了带有标记效果的而且单通道的图片了。

接下来我们选定一定数量的图片放到testing中,剩下的放到training中。注意:input里放原图,output放刚才生成的单通道的带有分割效果的图片。

然后我们在nnUNet-master/nnunet/dataset_conversion/Task120_Massachusetts_RoadSegm.py

这个py文件下找到几个我们要修改的变量,第一个base,base路径修改为我们刚才根目录下的road_segmentation_ideal文件的绝对路径,这里使用相对路径可能会报错,第二,修改task_name修改为我们要训练的项目名,前面提到的比如我用的就要改成Task01_Car。第三个要修改的是target_base,源代码用的是join函数,我这里直接改成了项目文件的路径,注意这里也是绝对路径!!比如我的

target_base = '/home/andre/桌面/nnUNetFrame_final/DATASET/nnUNet_raw/nnUNet_raw_data/Task01_Car'

继续往下看代码:看到这行代码:

training_cases = subfiles(labels_dir_tr, suffix='.jpg', join=False)

这里的suffix为我们input,output图片的路径,这里他默认的应该是png,我这边的数据都是jpg 的所以这里修改一下,这个py文件下有两个这段代码,所以要修改两次,找到suffix都改成jpg就ok。

然后运行这个py文件。运行成功后会在Task01_Car下生成五个文件,分别是imagesTr,imagesTs,labelsTr,labelsTs,dataset.json。

接下来运行指令:nnUNet_convert_decathlon_task -i /home/andre/桌面/nnUNetFrame_final/DATASET/nnUNet_raw/nnUNet_raw_data/Task01_Car

-i 后面跟的是Task01任务的路径 全文的每个路径都要根据自己的电脑进行修改,之后不再提示了。

运行成功后会在Task01_Car的同级目录下生成Task001_Car文件夹,注意观察新生成的imageTr文件中nii文件后四位是否是0000 0001 0002的形式如果末尾多出0000 则使用脚本删除多余的0000。

nnunet只识别末尾为0000.nii.gz 0001.nii.gz 0002.nii.gz格式的文件所以一定要检查好末尾。

要下图这样的格式,一个原图片变成了三个nii文件,末尾分别是0000 0001 0002

nnUNet使用教程(使用自己的数据训练自己的模型)_第2张图片

 这里再附一张json文件的格式:

nnUNet使用教程(使用自己的数据训练自己的模型)_第3张图片

做到这步我们接下来运行:nnUNet_plan_and_preprocess -t 1 

-t 1 再解释一下:-t后面跟的是你的任务号,Task01_Car的任务号为1

这步预处理成功之后会在preprocessed下生成Task001_Car文件夹

然后使用nnUNet_train 2d nnUNetTrainerV2 1 4进行训练 

简单解释一下这行代码

模型:2d

训练py文件:nnUNetTrainerV2

任务号:1

五折交叉验证:4

按照这个步骤我们只需要等待训练就ok了。

这是自己写的第一篇文章,希望大家多多支持。可能有些细节没有写到位,有不明白的可以问我。

找时间会继续更新batchsize的修改方法以及如何拿我们训练好的模型去验证,nnunet默认训练是1000轮时间比较长,我们可以中途拿bestmodel先测试一下。找时间会继续更新。


 

你可能感兴趣的:(nnunet,unet,深度学习,目标检测,人工智能,pytorch)