LITS 数据集 肝脏和肿瘤标签分离

标签分离原因说明

最近在用Unet做肝脏肿瘤分割，手中得到的数据集有LITS数据集，从网上下载下来的时候，格式是nii文件，因为我用的是2D Unet，所以要先把nii文件转为png格式。获取到的标签里面，肝脏和肿瘤是在一块的，所以想着把肝脏和肿瘤的标签分离出来再单独进行训练和预测，于是就有了这篇文章。

简单原理说明

在标签中，肝脏被标记为1，肿瘤被标记为2，所以，在加载图像后，对每一个图像中的不为1的像素置为0，以此得到只包含肝脏的标签。同理，把不为2的其他像素置为0，以此得到只包含肿瘤的标签。

代码实现

import os
import cv2
from medpy.io import load,save
import numpy as np

def label_seg(nii_path,save_folder):
    data_nii_path = nii_path + '/data'
    label_nii_path = nii_path+'/label'
    file_list = os.listdir(label_nii_path)    #批量处理
    for nii_label in file_list:

        liver_tumor_label, header = load(os.path.join(label_nii_path, nii_label))   # 加载nii文件
        liver_tumor_data, header1 = load(os.path.join(data_nii_path, nii_label.replace('segmentation', 'volume')))   # 加载nii文件
        liver_tumor_data[liver_tumor_data > 250] = 250	#加窗处理，过滤掉过于的组织
        liver_tumor_data[liver_tumor_data < -250] = -250
        liver = np.zeros(liver_tumor_label.shape)  # 创建一个全为0的矩阵，形状与加载的nii文件大小一样
        tumor = np.zeros(liver_tumor_label.shape)
        index = np.where(liver_tumor_label == 1)   # 肝脏标签为1，返回索引位置
        liver[index] = 1    # 将得到的肝脏坐标在新的矩阵中置为1，得到只包含肝脏的标签
        index1 = np.where(liver_tumor_label == 2)  # 肿瘤标签为2，返回索引位置
        tumor[index1] = 1   # 将得到的肿瘤坐标在新的矩阵中置为1，得到只包含肿瘤的标签
        slice_num = liver_tumor_label.shape[2]     # 获取切片数量
        tumor_path = save_folder+'/tumor_label'     # 分离的标签设置保存路径
        liver_path = save_folder+'/liver_label'
        train_tumor_path = save_folder+'/train_tumor'     # 训练数据保存路径
        train_liver_path = save_folder+'/train_liver'
        if not os.path.exists(tumor_path):      # 判断文件目录是否存在，不存在则创建
            os.mkdir(tumor_path)
        if not os.path.exists(liver_path):
            os.mkdir(liver_path)
        if not os.path.exists(train_tumor_path):      # 判断文件目录是否存在，不存在则创建
            os.mkdir(train_tumor_path)
        if not os.path.exists(train_liver_path):
            os.mkdir(train_liver_path)
        img_name = os.path.split(nii_label)     # 分离文件名，用以保存时的文件名前缀
        img_name = img_name[-1]
        img_name = img_name.split('.')
        img_name = img_name[0]
        for i in range(0, slice_num):

            if liver[:, :, i].max() > 0:    # 肝脏标签图片
                label_liver = liver[:, :, i] * 122
                liver_train_data = (liver_tumor_data[:, :, i] - liver_tumor_data[:, :, i].min()) / (liver_tumor_data[:, :, i].max() - liver_tumor_data[:, :, i].min()) * 255
                print(cv2.imwrite("%s/%s-%d.png" % (liver_path, "liver-"+img_name, i), label_liver))	
                print(cv2.imwrite("%s/%s-%d.png" % (train_liver_path, img_name.replace('segmentation', 'volume'), i), liver_train_data))	#保持分割出来的标签与数据的序号对应，因为肿瘤的标签量比较少
            if tumor[:, :, i].max() > 0:    # 肿瘤标签图片
                label_tumor = tumor[:, :, i] * 122
                tumor_train_data = (liver_tumor_data[:, :, i] - liver_tumor_data[:, :, i].min()) / (liver_tumor_data[:, :, i].max() - liver_tumor_data[:, :, i].min()) * 255
                print(cv2.imwrite("%s/%s-%d.png" % (tumor_path, "tumor-"+img_name, i), label_tumor))
                print(cv2.imwrite("%s/%s-%d.png" % (train_tumor_path, img_name.replace('segmentation', 'volume'), i),
                                  tumor_train_data))	#保持分割出来的标签与数据的序号对应，因为肿瘤的标签量比较少
    # save(liver, os.path.join(save_folder,"liver.nii"))
    # save(tumor, os.path.join(save_folder, "tumor.nii"))

if __name__ == "__main__":
    label_nii_path = r'nii 文件保存路径'

    save_path = 'png图像数据保存位置'
    label_seg(label_nii_path,save_path)

分离效果

结束语

本人水平有限，如有不当之处望请指正。