【AI达人特训营】基于PaddleSeg实现胃肠道MRI图像分割Baseline

Use PaddleSeg to finish the competition "UW-Madison GI Tract Image Segmentation".

AI Studio

556人浏览 · 2022-07-02 09:40:53

AI Studio · 2022-07-02 09:40:53 发布

基于PaddleSeg实现胃肠道MRI图像分割

在这里插入图片描述

本项目使用 PaddleSeg 完成kaggle的胃肠道MRI图像分割比赛。

任务:
在这个比赛中，你需要创建一个模型，自动分割胃肠道的核磁共振扫描。磁共振图像来自实际的癌症患者，他们在放疗期间分别做了1-5次磁共振成像扫描。你的算法将基于这些扫描数据，提出创造性的深度学习解决方案，帮助癌症患者获得更好的治疗。

PaddleSeg:
PaddleSeg是基于飞桨PaddlePaddle开发的端到端图像分割开发套件，涵盖了高精度和轻量级等不同方向的大量高质量分割模型。通过模块化的设计，提供了配置化驱动和API调用两种应用方式，帮助开发者更便捷地完成从训练到部署的全流程图像分割应用。

除了数据预处理之外，使用PaddleSeg可以节省大量时间。本项目提供了预处理的代码，因此您基本上可以在没有额外代码的情况下完成比赛。如果这个项目帮你节省了时间，可以帮忙点下fork和star。

1. 加载库

安装并导入所使用的库。

!pip install -qq pandarallel

!pip install -qq collection

import os
import cv2
import random
import imageio
import numpy as np
import pandas as pd
from tqdm import tqdm
from glob import glob
from PIL import Image
from collections import Counter
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

from pathlib import Path
from IPython.display import IFrame
from sklearn.model_selection import train_test_split

from joblib import Parallel, delayed
from pandarallel import pandarallel
# Initialization(设置为CPU核心数，达到最佳效率)
pandarallel.initialize(nb_workers=4)

2. 预处理 (可以跳过，直接训练模型)

这部分将使用原始数据集来生成一个可以被PaddleSeg直接使用的数据集。如果不需要更改数据集，可以跳过本节。

Resources:

https://en.wikipedia.org/wiki/Run-length_encoding
https://www.kaggle.com/paulorzp/run-length-encode-and-decode
https://www.kaggle.com/code/awsaf49/uwmgi-mask-data
https://github.com/PaddlePaddle/PaddleSeg/blob/release/2.5/docs/data/marker/marker_cn.md

2.1 载入原始数据

确保你加载了这个数据集:
https://aistudio.baidu.com/aistudio/datasetdetail/145482

Resource:

# 解压数据集
!unzip -qoa data/data145482/uw-madison-gi-tract-image-segmentation.zip -d data

# 定义像素值和类对应关系
class2pixel_value_dic = {'background':0, 'large_bowel':1, 'small_bowel':2, 'stomach':3}

# 读入csv
df = pd.read_csv('data/train.csv')

# 利用id生成case, day和slice
def get_metadata(row):
    data = row['id'].split('_')
    case = int(data[0].replace('case',''))
    day = int(data[1].replace('day',''))
    slice_ = int(data[-1])
    row['case'] = case
    row['day'] = day
    row['slice'] = slice_
    return row
    
# 将文件名转换成有效信息
def path2info(row):
    path = row['image_path']
    data = path.split('/')
    slice_ = int(data[-1].split('_')[1])
    case = int(data[-3].split('_')[0].replace('case',''))
    day = int(data[-3].split('_')[1].replace('day',''))
    width = int(data[-1].split('_')[2])
    height = int(data[-1].split('_')[3])
    row['height'] = height
    row['width'] = width
    row['case'] = case
    row['day'] = day
    row['slice'] = slice_
    return row

# Parallel apply(pandas apply的并行执行)
# 利用id生成case, day和slice
df = df.parallel_apply(get_metadata,axis=1)

# 将文件名转换成有效信息
paths = glob('data/train/*/*/*/*')
path_df = pd.DataFrame(paths, columns=['image_path'])
path_df = path_df.parallel_apply(path2info, axis=1)

df_merge = pd.merge(df, path_df, on=['case','day','slice'])

def load_img(path):
    '''
    根据path载入图片
    '''
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    img = img.astype('float32') # original is uint16
    img = (img - img.min())/(img.max() - img.min())*255.0 # scale image to [0, 255]
    img = img.astype('uint8')
    return img

def show_img(img, mask=None):
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    img = clahe.apply(img)
    plt.imshow(img, cmap='bone')
    if mask is not None:
        plt.imshow(mask, alpha=0.5)
        handles = [Rectangle((0,0),1,1, color=_c) for _c in [(0.667,0.0,0.0), (0.0,0.667,0.0), (0.0,0.0,0.667)]]
        labels = [ "Large Bowel", "Small Bowel", "Stomach"]
        plt.legend(handles,labels)
    plt.axis('off')

# RLE
# https://en.wikipedia.org/wiki/Run-length_encoding
# https://www.kaggle.com/paulorzp/run-length-encode-and-decode
def rle_decode(mask_rle, shape, pixel_value, img):
    '''
    mask_rle: csv中的原始RLE
    shape: (height,width) 返回的shape 
    pixel_value: mask中的像素值(取决于类)
    img: 正在等待添加新区域的不完整mask

    Return: numpy array
    '''
    s = np.asarray(mask_rle.split(), dtype=int)
    starts = s[0::2] - 1
    lengths = s[1::2]
    ends = starts + lengths
    img = img.reshape(-1)
    for lo, hi in zip(starts, ends):
        img[lo:hi] = pixel_value
    return img.reshape(shape)  # 需要对齐到图片的shape

def id2mask(id_, df):
    '''
    input: 
        id_: 图片的id
        df: (dataFrame) 需要有id, height, weight, class, segmentation
    
    output:
        mask: 单通道mask
    '''
    idf = df[df['id']==id_]
    wh = idf[['height','width']].iloc[0]
    shape = (wh.height, wh.width)
    mask = np.zeros(shape, dtype=np.uint8)

    # mask = np.zeros(shape, dtype=np.uint8)
    for i, class_ in enumerate(['large_bowel', 'small_bowel', 'stomach']):
        cdf = idf[idf['class']==class_]
        rle = cdf.segmentation.squeeze()
        pixel_value = class2pixel_value_dic[class_]
        if len(cdf) and not pd.isna(rle):
            mask = rle_decode(rle, shape, pixel_value, mask)
    return mask

def mask_to_onehot(mask, palette):
    """
    将mask从(H,W)/(H, W, C) 转换为(H, W, K), C通常是1或3, K是类别的数量。
    palette是一个数组，包含掩码中的像素值(可以没有背景)
    """
    if mask.ndim==2:
        mask = np.expand_dims(mask, axis=-1)

    semantic_map = []
    for colour in palette:
        equality = np.equal(mask, colour)
        class_map = np.all(equality, axis=-1)
        semantic_map.append(class_map)
    semantic_map = np.stack(semantic_map, axis=-1).astype(np.float32)
    return semantic_map

def onehot_to_mask(mask, palette):
    """
    将mask从(H, W, K)转化为(H, W, C)
    """
    x = np.argmax(mask, axis=-1)
    colour_codes = np.array(palette)
    x = np.uint8(colour_codes[x.astype(np.uint8)])
    return x

# 看看mask效果
palette = [1, 2, 3]  #without background
img = load_img(df_merge[df_merge['id']=='case123_day20_slice_0082'].image_path.iloc[0])
gt_onehot = mask_to_onehot(id2mask('case123_day20_slice_0082', df_merge), palette)
show_img(img, gt_onehot)

在这里插入图片描述

2.2 保存mask为图片

PaddleSeg使用单通道的mask图像，每个像素值代表一个类别，像素标签类别需要从0增加。例如，0、1、2、3表示有4个类别。

Pixel Value	Class
0	Background
1	Large bowel
2	Small bowel
3	Stomach

def save_mask(id_, df=None):
    '''
    input:
        id_:  想要保存mask的原图片的id
        df: (dataFrame)需要包含id, height, weight, class, segmentation, scan path...

    output:
        保存的mask文件的路径(把scans图片路径中'scans'替换为'masks')
    '''
    idf = df[df['id']==id_]
    mask = id2mask(id_, df=df) # one channel

    image_path = idf.image_path.iloc[0]
    mask_path = image_path.replace('scans','masks')
    mask_folder = mask_path.rsplit('/',1)[0]
    os.makedirs(mask_folder, exist_ok=True)
    cv2.imwrite(mask_path, mask)

    return mask_path

# 看看save_mask函数
_ = save_mask('case123_day20_slice_0082', df_merge)

palette = [1, 2, 3]  #without background
img = load_img(df_merge[df_merge['id']=='case123_day20_slice_0082'].image_path.iloc[0])
mask_path = df_merge[df_merge['id']=='case123_day20_slice_0082'].image_path.iloc[0].replace('scans','masks')
mask = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
gt_onehot = mask_to_onehot(mask, palette)
show_img(img, gt_onehot)

在这里插入图片描述

# 找到所有有mask的MRI图像
tmp_df = df_merge.copy()
tmp_df = tmp_df[~df.segmentation.isna()]
ids = tmp_df['id'].unique()
# 保存所有mask
_ = Parallel(n_jobs=-1, backend='threading')(delayed(save_mask)(id_, df=tmp_df)\
                                             for i, id_ in enumerate(tqdm(ids, total=len(ids))))

100%|██████████| 16590/16590 [02:09<00:00, 125.20it/s]

# 更新df_merge中的mask信息
df_merge['mask_path']=np.nan
for i in range(0,len(df_merge),3):
    # 如果不是所有类别的'segmentation'都是NaN，图像应该有mask_path
    if not all(df_merge['segmentation'][i:i+3].isna()):
        df_merge['mask_path'][i:i+3]=df_merge['image_path'][i].replace('scans','masks')

2.3 统一图片和mask大小

# 获取所有scans和masks的路径
allScansPaths = glob('data/train/*/*/scans/*')
allMasksPaths = glob('data/train/*/*/masks/*')

# 对于scans, 使用双线性插值
for imagePath in allScansPaths:
    image = cv2.imread(imagePath, cv2.IMREAD_UNCHANGED)
    image = cv2.resize(image, (256, 256))     # to 256*256 by INTER_LINEAR
    _ = os.remove(imagePath)                  # delete original image
    _ = cv2.imwrite(imagePath, image)         # save resized image

# 对于masks, 使用最近邻插值
for imagePath in allMasksPaths:
    image = cv2.imread(imagePath, cv2.IMREAD_UNCHANGED)
    image = cv2.resize(image, (256, 256), interpolation=cv2.INTER_NEAREST)     # to 256*256 by INTER_NEAREST
    _ = os.remove(imagePath)                                                   # delete original image
    _ = cv2.imwrite(imagePath, image)                                          # save resized image

# 更改高度和宽度信息
def fix_size4metadata(row, height, width):
    row['height'] = height
    row['width'] = width
    return row

# 更改metadata中的高度和宽度信息
df_merge = df_merge.apply(fix_size4metadata,axis=1,args=(256,256,))

# 看看resize后的图片
palette = [1, 2, 3]  #without background
img = load_img(df_merge[df_merge['id']=='case123_day20_slice_0082'].image_path.iloc[0])
mask_path = df_merge[df_merge['id']=='case123_day20_slice_0082'].image_path.iloc[0].replace('scans','masks')
mask = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
gt_onehot = mask_to_onehot(mask, palette)
show_img(img, gt_onehot)

在这里插入图片描述

2.4 分割数据集

在本小节中，我们将分层的把95%的数据分成训练集，其余部分分成验证集。这个比例不是固定的，你可以根据自己的意愿改变它。

# 为每个x_id设置三个one-hot标签，以表明mask包含哪些器官
df_merge['large_bowel']=np.nan
df_merge['small_bowel']=np.nan
df_merge['stomach']=np.nan
for i in range(0,len(df_merge),3):
    # 三个one-hot标签，标明mask包含哪些器官
    df_merge['large_bowel'][i:i+3]= 0 if pd.isnull(df_merge['segmentation'][i]) else 1
    df_merge['small_bowel'][i:i+3]= 0 if pd.isnull(df_merge['segmentation'][i+1]) else 1
    df_merge['stomach'][i:i+3]= 0 if pd.isnull(df_merge['segmentation'][i+2]) else 1

# 获取所有具有mask的图像id
x_id = df_merge[~pd.isnull(df_merge['segmentation'])]['id'].unique()

# 分配内存
y_onehot = np.array([[0.0,0.0,0.0]]*len(x_id))
# 为每个x_id设置三个one-hot标签，以表明mask包含哪些器官
for i in range(len(x_id)):
    y_onehot[i] = np.array(df_merge[df_merge['id']==x_id[i]][['large_bowel', 'small_bowel', 'stomach']].iloc[0])

# 看看数据集的形状 
x_id.shape, y_onehot.shape

((16590,), (16590, 3))

# 分层分割数据集
X_train_id, X_val_id,  y_train_onehot, y_val_onehot = train_test_split(x_id, y_onehot, test_size = 0.05, stratify=y_onehot, random_state=42)

# 看看分割后是不是平衡的
pd.DataFrame({
    'train': Counter(str(row) for row in y_train_onehot),
    'test' : Counter(str(row) for row in y_val_onehot)
}).T.fillna(0.0)

	[1. 1. 1.]	[1. 0. 1.]	[1. 1. 0.]	[0. 0. 1.]	[0. 1. 1.]	[1. 0. 0.]	[0. 1. 0.]
train	3041	2831	7392	2171	152	117	56
test	160	149	389	115	8	6	3

2.5 保存预处理后的数据集

拆分数据集后，我们需要生成train.txt和val.txt文件。
https://github.com/PaddlePaddle/PaddleSeg/blob/release/2.5/docs/data/marker/marker_cn.md

# 生成 labels txt
with open("data/train/labels.txt","w") as f:
    f.write("background\n")
    f.write("large_bowel\n")
    f.write("small_bowel\n")
    f.write("stomach")

# 生成 train txt
with open("data/train/train.txt","w") as f:
    for id_ in X_train_id:
        imagePath = '/'.join((df_merge[df_merge['id']==id_]['image_path'].iloc[0]).split('/')[2:])
        maskPath  =  imagePath.replace('scans', 'masks')
        f.write(imagePath)
        f.write(" ")
        f.write(maskPath)
        f.write("\n")

# 生成 validation txt
with open("data/train/val.txt","w") as f:
    for id_ in X_val_id:
        imagePath = '/'.join((df_merge[df_merge['id']==id_]['image_path'].iloc[0]).split('/')[2:])
        maskPath  =  imagePath.replace('scans', 'masks')
        f.write(imagePath)
        f.write(" ")
        f.write(maskPath)
        f.write("\n")

# 最终metadata的样子
df_merge.head()

	id	class	segmentation	case	day	slice	image_path	height	width	mask_path
0	case123_day20_slice_0001	large_bowel	NaN	123	20	1	data/train/case123/case123_day20/scans/slice_0...	256	256	NaN
1	case123_day20_slice_0001	small_bowel	NaN	123	20	1	data/train/case123/case123_day20/scans/slice_0...	256	256	NaN
2	case123_day20_slice_0001	stomach	NaN	123	20	1	data/train/case123/case123_day20/scans/slice_0...	256	256	NaN
3	case123_day20_slice_0002	large_bowel	NaN	123	20	2	data/train/case123/case123_day20/scans/slice_0...	256	256	NaN
4	case123_day20_slice_0002	small_bowel	NaN	123	20	2	data/train/case123/case123_day20/scans/slice_0...	256	256	NaN

# 保存df_merge，其中包含所有必要的信息
df_merge.to_csv('data/train_metadata.csv', index=False)

最后，把所有东西打包成zip文件。
并上传到Aistudio Dataset, 避免每次都进行预处理。

3. 加载数据集 (加载预处理过的数据集)

跳过预处理时，需要使用本节。
确保您挂载了经过预处理的数据集，对于这个notebook，预处理后的数据集是https://aistudio.baidu.com/aistudio/datasetdetail/153829 .

与原始数据集不同的地方:

根据原始数据将RLE的mask转换为png，mask保存路径为将图像路径中“scans”替换为“masks”。
将所有scans和masks转换为256*256
包含"train.txt"和"val.txt"在"train"文件夹中。(为使用PaddleSeg)
train_metadata.csv包含所有必要信息。

# 解压缩数据集，其中包括所有必要的文件
!unzip -qoa data/data153829/Processed_gi-tract-image-segmentation.zip -d data

4. PaddleSeg

使用PaddleSeg，您可以通过简单地修改yml文件加上两行命令来完成训练和推断过程。
如前所述，这个项目只是一个例子，您可以参考官方文档对训练过程做出改变，让模型有更好的性能。您可以更改包括但不限于models, losses 和 augmentations.

# clone PaddleSeg
# 加速clone(China)
!git clone https://gitee.com/paddlepaddle/PaddleSeg.git 
# General
# !git clone https://github.com/PaddlePaddle/PaddleSeg.git

# 安装 PaddleSeg
!pip install -qq paddleseg

fatal: 目标路径 'PaddleSeg' 已经存在，并且不是一个空目录。

# 切换jupyter工作目录!
%cd PaddleSeg

/home/aistudio/PaddleSeg

# !pip install -r requirements.txt

4.1 训练

在定义完yml文件之后，使用以下命令开始训练。Checkpoints和最佳模型保存在您定义的—save_dir目录中。同时，您可以使用VisualDL来可视化存储在“–save_dir”目录中的日志文件。

需要注意的一点是，RandomPaddingCrop的默认填充值是255。因为我们的像素标签类别需要从0增加，mask上的填充将影响dice loss的使用。因此，您需要在yml文件中的RandomPaddingCrop下添加一个额外的参数“label_padding_value: 0”。

一些有用的yml文件作为参考:

!python train.py \
       --config ../unet_256_256_woTL.yml \
       --do_eval \
       --use_vdl \
       --num_workers 4 \
       --save_interval 2000 \
       --save_dir output_unet_256_256_woTL

4.2 预测

除了分析IOU, ACC和Kappa，我们还可以检查一些具体样本的分割效果，并从Bad Case中获得进一步的优化思路。

在进行预测时，——image_path可以是图片的路径、包含图像路径的文件列表或目录。此时，将对该图像或文件列表或目录中的所有图像进行预测，并保存可视化结果。

Resource:
https://github.com/PaddlePaddle/PaddleSeg/blob/release/2.5/docs/predict/predict_cn.md

df_merge = pd.read_csv('../data/train_metadata.csv')

def load_img(path):
    '''
    加载图片
    '''
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    img = img.astype('float32') # original is uint16
    img = (img - img.min())/(img.max() - img.min())*255.0 # scale image to [0, 255]
    img = img.astype('uint8')
    return img

def show_img(img, mask=None):
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    img = clahe.apply(img)
    plt.imshow(img, cmap='bone')
    if mask is not None:
        plt.imshow(mask, alpha=0.5)
        handles = [Rectangle((0,0),1,1, color=_c) for _c in [(0.667,0.0,0.0), (0.0,0.667,0.0), (0.0,0.0,0.667)]]
        labels = [ "Large Bowel", "Small Bowel", "Stomach"]
        plt.legend(handles,labels)
    plt.axis('off')

def id2mask(id_, df):
    '''
    input: 
        id_: 图片id
        df: (dataFrame)需要包含id, height, weight, class, segmentation
    
    output:
        mask: 单通道mask
    '''
    idf = df[df['id']==id_]
    wh = idf[['height','width']].iloc[0]
    shape = (wh.height, wh.width)
    mask = np.zeros(shape, dtype=np.uint8)

    # mask = np.zeros(shape, dtype=np.uint8)
    for i, class_ in enumerate(['large_bowel', 'small_bowel', 'stomach']):
        cdf = idf[idf['class']==class_]
        rle = cdf.segmentation.squeeze()
        pixel_value = class2pixel_value_dic[class_]
        if len(cdf) and not pd.isna(rle):
            mask = rle_decode(rle, shape, pixel_value, mask)
    return mask

def mask_to_onehot(mask, palette):
    """
    将mask从(H,W)/(H, W, C) 转换为(H, W, K), C通常是1或3, K是类别的数量。
    palette是一个数组，包含掩码中的像素值(可以没有背景)
    """
    if mask.ndim==2:
        mask = np.expand_dims(mask, axis=-1)

    semantic_map = []
    for colour in palette:
        equality = np.equal(mask, colour)
        class_map = np.all(equality, axis=-1)
        semantic_map.append(class_map)
    semantic_map = np.stack(semantic_map, axis=-1).astype(np.float32)
    return semantic_map

# 预测
# 从验证集中随机选择一张图片作为例子
!python predict.py \
       --config ../unet_256_256_woTL.yml \
       --model_path output_unet_256_256_woTL/best_model/model.pdparams \
       --image_path ../data/train/case20/case20_day0/scans/slice_0120_266_266_1.50_1.50.png \
       --save_dir output_unet_256_256_woTL/result \
       --custom_color 0 0 0 0 0 255 0 255 0 255 0 0

# 真实的图片和mask
palette = [1, 2, 3]  #without background
img = load_img('../data/train/case20/case20_day0/scans/slice_0120_266_266_1.50_1.50.png')
gt_onehot = mask_to_onehot(cv2.imread('../data/train/case20/case20_day0/masks/slice_0120_266_266_1.50_1.50.png'), palette)
show_img(img, gt_onehot)

在这里插入图片描述

# 预测的图片和mask
plt.imshow(cv2.imread('output_unet_256_256_woTL/result/added_prediction/slice_0120_266_266_1.50_1.50.png'))

<matplotlib.image.AxesImage at 0x7f0722ff1510>

在这里插入图片描述

5. 总结

从整个notebook中很容易看出，大多数时间都花在处理数据上，paddleseg极大地简化了训练和推理过程。在处理大量数据时，请不要忘记使用一些支持并行的库来加快流程，例如joblib和pandarallel。

当我们使用处理过的数据时，我们可以非常高效地训练一个模型，几乎不需要任何代码。使用VisualDL将我们的训练过程可视化，我们可以发现我们的模型的性能可以通过训练更多的时间来提高。然而，更有效的提高性能的方法可以是使用迁移学习，一个更好的augmentation策略，更适合分割的loss。

在这里插入图片描述