飞桨常规赛：视杯视盘分割（GAMMA挑战赛任务三） - 11月第5名方案

MICCAI2021 Contest : GAMMA 任务三：对2D眼底图像中的视杯和视盘结构进行分割

AI Studio

601人浏览 · 2022-01-07 22:34:52

AI Studio · 2022-01-07 22:34:52 发布

飞桨常规赛：视杯视盘分割（GAMMA挑战赛任务三） - 11月第5名方案

〇、数据解压及项目准备

# 数据集解压
! mkdir -p datasets
! unzip -oq work/Disc_Cup.zip -d datasets
# 解压测试数据
! mkdir -p tests
! unzip -oq work/test.zip -d tests

# 项目准备
! git clone --depth=1 https://gitee.com/paddlepaddle/PaddleSeg.git  # paddleseg，github太慢
! pip install -q patta  # patta
# ! pip install -q albumentations  # 这个可以在终端通过pip3安装

import sys

sys.path.append('PaddleSeg')  # paddleseg

一、数据准备

1.1 数据信息查看

通过图像信息的查看，我们可以有两种不同大小的图像，因此在后续数据增强的时候我们可以对这两个大小的数据进行crop、padding和resize等手段，使他们大小相同且眼睛的大小和位置也相同。

# import os
# import numpy as np
# from PIL import Image

# 查看大小
# img_size = []
# imgs_folder_path = 'datasets/Disc_Cup_Mask'
# imgs_name = os.listdir(imgs_folder_path)
# for name in imgs_name:
#     img_path = os.path.join(imgs_folder_path, name)
#     img = np.asarray(Image.open(img_path))
#     img_size.append(img.shape)
# print(set(img_size))
# 查看标签数值
# label = np.asarray(Image.open('datasets/Disc_Cup_Mask/0004.png'))
# print(set(label.flatten()))

{(1934, 1956), (2000, 2992)}
{0, 128, 255}

1.2 划分数据集

由于数据较少，担心随机选择的验证集过于相似无法更好的评估结果，所以手动选择了10张差异较大的图像用于评估。

# import os
# # import random
# from PIL import Image

# # 手动选择了10张特征各异的，评估比较有说服力
# val_name_list = ['0080.jpg', '0070.jpg', '0006.jpg', '0063.jpg', '0086.jpg', \
#                  '0075.jpg', '0096.jpg', '0030.jpg', '0062.jpg', '0081.jpg']

# def create_list(data_path):
#     image_path = os.path.join(data_path, 'Image')
#     label_path = os.path.join(data_path, 'Disc_Cup_Mask')
#     data_names = os.listdir(image_path)
#     # random.shuffle(data_names)  # 打乱数据
#     with open(os.path.join(data_path, 'train_list.txt'), 'w') as tf:
#         with open(os.path.join(data_path, 'val_list.txt'), 'w') as vf:
#             for idx, data_name in enumerate(data_names):
#                 img = os.path.join('Image', data_name)
#                 lab = os.path.join('Disc_Cup_Mask', data_name.replace('jpg', 'png'))
#                 # if idx % 9 == 0:  # 90%的作为训练集
#                 if data_name in val_name_list:
#                     vf.write(img + ' ' + lab + '\n')
#                 else:
#                     tf.write(img + ' ' + lab + '\n')
#     print('数据列表生成完成')

# data_path = 'datasets'
# create_list(data_path)  # 生成数据列表

1.3 构建数据集

其中三个新建的transform需要在源代码中查看：

NormSize：会将两种大小的图像统一成一样的，同时对于矩形图像两边过多的黑色背景，对图像的范围进行了选择，使得眼睛区域尽可能充满整个图像。
LocalEqualHist：会对图像进行自适应局部直方图均衡化，使图像的特征看起来更明显。
HighlightClipping：根据最亮点像素裁剪512x512的区域（由于查看大多数眼睛图像的最亮点都在视盘区中心）。但有少数不是，且使用后效果没有提高，故没有使用了。

import paddleseg.transforms as T
from paddleseg.datasets import Dataset

# 构建训练集
train_transforms = [
    T.NormSize(),  # 自定义标准化图像大小2000x2000
    T.LocalEqualHist(),  # 自适应局部直方图均衡化
    T.RandomHorizontalFlip(),  # 水平翻转
    T.RandomRotation(im_padding_value=[0, 0, 0]),  # 随机旋转
    T.Resize(target_size=(1024, 1024)),  # 兼顾大小
    # T.HighlightClipping(),
    T.Normalize(
        [0.3883131, 0.25449154, 0.110598095], 
        [0.26274413, 0.1827712, 0.12587263]),  # 标准化
]
train_dataset = Dataset(
    transforms=train_transforms,
    dataset_root='datasets',
    num_classes=3,
    mode='train',
    train_path='datasets/train_list.txt',
    separator=' ',
)
# 构建验证集
val_transforms = [
    T.NormSize(),
    T.LocalEqualHist(),
    T.Resize(target_size=(1024, 1024)),
    # T.HighlightClipping(),
    T.Normalize(
        [0.3883131, 0.25449154, 0.110598095], 
        [0.26274413, 0.1827712, 0.12587263]),    
]
val_dataset = Dataset(
    transforms=val_transforms,
    dataset_root='datasets',
    num_classes=3,
    mode='train',  # 这里用train，不然无法对标签进行NormSize
    train_path='datasets/val_list.txt',
    separator=' ',
)

# 测试输出
import matplotlib.pyplot as plt

%matplotlib inline

for img, lab in train_dataset:
    print(set(lab.flatten()))
    print(img.shape, lab.shape)
    plt.subplot(121);plt.imshow(img.transpose((1, 2, 0)).astype('uint8'))
    plt.subplot(122);plt.imshow(lab)
    plt.show()
    break

1.4 查看增强数据

*ps：查看的时候需要把上述的Normalize注释掉，不然无法看出原图像。

# ! mkdir -p vis_aug
# ! mkdir -p vis_aug/img
# ! mkdir -p vis_aug/lab

# from PIL import Image
# from tqdm import tqdm

# for idx, (img, lab) in enumerate(tqdm(train_dataset)):
#     pimg = Image.fromarray(img.transpose(1, 2, 0).astype('uint8'))
#     plab = Image.fromarray(lab.astype('uint8'))
#     pimg.save('vis_aug/img/vis_' + str(idx) + '.jpg')
#     plab.save('vis_aug/lab/vis_' + str(idx) + '.png')

# ! mkdir -p vis_aug_test
# ! mkdir -p vis_aug_test/img
# ! mkdir -p vis_aug_test/lab

# from PIL import Image
# from tqdm import tqdm

# for idx, (img, lab) in enumerate(tqdm(val_dataset)):
#     pimg = Image.fromarray(img.transpose(1, 2, 0).astype('uint8'))
#     plab = Image.fromarray(lab.astype('uint8'))
#     pimg.save('vis_aug_test/img/vis_' + str(idx) + '.jpg')
#     plab.save('vis_aug_test/lab/vis_' + str(idx) + '.png')

1.5 均值、方差统计

统计图像的均值方差后，回到前面填入Normalize中。

# import os
# import cv2
# import numpy as np
# from tqdm import tqdm
 
# means, stdevs = [], []
# img_list = []
# imgs_path = 'vis_aug/img'
# imgs_name = os.listdir(imgs_path)
# for name in tqdm(imgs_name):
#     img = cv2.cvtColor(cv2.imread(os.path.join(imgs_path, name)), cv2.COLOR_BGR2RGB)
#     img = img[:, :, :, np.newaxis]
#     img_list.append(img)
# imgs = np.concatenate(img_list, axis=-1)
# imgs = imgs.astype(np.float32) / 255.
# for i in range(3):
#     pixels = imgs[:, :, i, :].ravel()  # 拉成一行
#     means.append(np.mean(pixels))
#     stdevs.append(np.std(pixels))
# print(means, stdevs)

[0.3883131, 0.25449154, 0.110598095] [0.26274413, 0.1827712, 0.12587263]

二、训练准备

import paddle
import paddle.nn as nn
from paddleseg.models import OCRNet, HRNet_W18

model = model = OCRNet(
    backbone=HRNet_W18(),
    backbone_indices=[0],
    num_classes=3)
params = paddle.load('output_ocrnet_hrnet18/best_model/model.pdparams')
model.set_state_dict(params)

from paddleseg.models.losses import CrossEntropyLoss, DiceLoss, MixedLoss

iters = 5000
batch_size = 4
base_lr = 3e-4

lr = paddle.optimizer.lr.CosineAnnealingDecay(base_lr, T_max=int(iters // 1.5))
optimizer = paddle.optimizer.Adam(
    learning_rate=lr,
    parameters=model.parameters(),
    weight_decay=paddle.regularizer.L2Decay(1e-7),
    grad_clip=paddle.nn.ClipGradByGlobalNorm(clip_norm=1.0))
losses = {}
losses['types'] = [MixedLoss([CrossEntropyLoss(), DiceLoss()], [1, 1])] * 5  # 2
losses['coef'] = [1] * 5  # 2

三、模型训练

from paddleseg.core import train

train(
    model=model,
    train_dataset=train_dataset,
    val_dataset=val_dataset,
    optimizer=optimizer,
    save_dir='output_ocrnet_hrnet18',
    iters=iters,
    batch_size=batch_size,
    save_interval=int(iters/10),
    log_iters=10,
    num_workers=0,
    losses=losses,
    use_vdl=True)

四、模型预测

使用了水平翻转的tta，由于图像较大，使用较多的tta会显存不够。
需要将大小和位置统一回原数据，这里和NormSize一样提前算好了区域，所以才能对应还原到原来大小。

import paddle
from paddleseg.models import OCRNet, HRNet_W18
import paddleseg.transforms as T
from paddleseg.core import infer
import os
from tqdm import tqdm
from PIL import Image
import numpy as np
import patta as tta
from mytools import tensor2result, restore

def nn_infer(model, imgs_path, is_tta=True):
    if not os.path.exists('result'):
        os.mkdir('result')
    # 预测结果
    transforms = T.Compose([
        T.NormSize(),
        T.LocalEqualHist(),
        T.Resize(target_size=(1024, 1024)),
        T.Normalize(
            [0.3883131, 0.25449154, 0.110598095], 
            [0.26274413, 0.1827712, 0.12587263])
    ])
    # 循环预测和保存
    for img_path in tqdm(imgs_path):
        H, W = np.asarray(Image.open(img_path)).shape[:2]  # 获取原始的H和W
        img, _ = transforms(img_path)  # 进行数据预处理
        img = paddle.to_tensor(img[np.newaxis, :])  # C,H,W -> 1,C,H,W
        # TTA
        if is_tta == True:
            tta_pres = paddle.zeros([1, 3, 1024, 1024])  # 图像大小1024
            for tta_transform in tta.aliases.hflip_transform ():
                tta_img = tta_transform.augment_image(img)  # TTA_transforms
                tta_pre = infer.inference(model, tta_img)  # 预测
                deaug_pre = tta_transform.deaugment_mask(tta_pre)
                tta_pres += deaug_pre
            pre = tta_pres / 2.
        else:
            pre = infer.inference(model, img)  # 预测
        pred = tensor2result(pre)  # 转为颜色对应的array
        pred = restore(pred, H, W)  # 恢复原始大小
        pil_img = Image.fromarray(pred)
        pil_img.save(os.path.join('result', img_path.split('/')[-1].replace('jpg', 'bmp')), 'bmp')

# 网络准备
model_path = 'output_ocrnet_hrnet18/best_model/model.pdparams'
model = OCRNet(
    backbone=HRNet_W18(),
    backbone_indices=[0],
    num_classes=3)
params = paddle.load(model_path)
model.set_state_dict(params)
model.eval()
# 预测文件
# set_path = 'datasets'
# list_file = 'datasets/val_list.txt'
# imgs_path = []
# with open(list_file, 'r') as f:
#     datas_path = f.readlines()
#     for data_path in datas_path:
#         imgs_path.append(os.path.join(set_path, data_path.split(' ')[0].strip()))
test_folder = "tests"
imgs_path = os.listdir(test_folder)
imgs_path = [os.path.join(test_folder, name) for name in imgs_path]
# 预测
nn_infer(model, imgs_path, is_tta=True)

100%|██████████| 100/100 [01:05<00:00,  1.52it/s]

五、预测后处理

由于结果只是一个实心的区域，因此可以使用闭运算填充孔洞。
保留最大联通区，去掉其他小的联通区。

import os
import numpy as np
from tqdm import tqdm
from PIL import Image
from aftercure import one_package_service

pred_folder = 'result'
save_folder = 'af_result'
if not os.path.exists(save_folder):
    os.mkdir(save_folder)
imgs_name = os.listdir(pred_folder)
for name in tqdm(imgs_name):
    img_path = os.path.join(pred_folder, name)
    img = np.asarray(Image.open(img_path))
    result = Image.fromarray(one_package_service(img, k_size=300))
    result.save(img_path.replace(pred_folder, save_folder))

100%|██████████| 100/100 [00:18<00:00,  5.54it/s]

六、结果评价

Dice系数作为分割结果测评的评价指标：

$\text { Dice }=\frac{2|X \cap Y|}{|X|+|Y|}$

其中，X代表金标准中分割目标像素点集合，Y代表预测结果中分割目标像素点集合，公式中 |X∩Y| 是X和Y之间的交集，|X|和|Y|分表表示X和Y的元素的个数。
此外，我们使用平均绝对误差（Mean Absolute Error）来测量样本视杯视盘分割结果与金标准之间的垂直杯盘比差异。垂直杯盘比有直接的临床相关性，可辅助评估青光眼的进展。求解方式为计算垂直方向上视杯区域和视盘区域最大直径的比值。对于任务三，最终的评分综合了视杯分割结果Dice系数、视盘分割结果Dice系数，和垂直杯盘比数值的MAE：

$\text { Score }_{\text {task } 3}=0.25 * \text { Dice }_{\text {cup }}+0.35 * \text { Dice }_{d i s c}+0.04 * \frac{1}{M A E+0.1}$

*ps：貌似写的有点问题，和官方的得分有所区别。

# from mytools import get_score

# score = get_score('af_result', 'datasets/Disc_Cup_Mask', False)
# print('当前得分为:', score)

当前得分为: 0.7742816652233033

显示结果

import matplotlib.pyplot as plt
from PIL import Image
import numpy as np

%matplotlib inline

img = np.asarray(Image.open("tests/0113.jpg"))
pre = np.asarray(Image.open("result/0133.bmp"))
end = np.asarray(Image.open("af_result/0133.bmp"))
plt.figure(figsize=(15, 5))
plt.subplot(131);plt.imshow(img);plt.title("img")
plt.subplot(132);plt.imshow(pre);plt.title("pre")
= np.asarray(Image.open("tests/0113.jpg"))
pre = np.asarray(Image.open("result/0133.bmp"))
end = np.asarray(Image.open("af_result/0133.bmp"))
plt.figure(figsize=(15, 5))
plt.subplot(131);plt.imshow(img);plt.title("img")
plt.subplot(132);plt.imshow(pre);plt.title("pre")
plt.subplot(133);plt.imshow(end);plt.title("end")