[论文解读] Two layer Ensemble of Deep Learning Models for Medical Image Segmentation

Truong Dang, Tien Thanh Nguyen+3 more|arXiv (Cornell University)|Apr 10, 2021

Advanced Neural Network Applications参考文献 30被引用 11

一句话总结

该论文提出了一种用于医学图像分割的两层集成深度学习框架，其中第一层模型的预测结果被用作第二层模型的增强输入。最终输出通过线性回归结合学习到的权重生成，在CAMUS和Kvasir-SEG数据集上实现了最先进性能，Dice分数最高达0.96，Hausdorff距离相比基线模型最高改善0.3。

ABSTRACT

In recent years, deep learning has rapidly become a method of choice for the segmentation of medical images. Deep Neural Network (DNN) architectures such as UNet have achieved state-of-the-art results on many medical datasets. To further improve the performance in the segmentation task, we develop an ensemble system which combines various deep learning architectures. We propose a two-layer ensemble of deep learning models for the segmentation of medical images. The prediction for each training image pixel made by each model in the first layer is used as the augmented data of the training image for the second layer of the ensemble. The prediction of the second layer is then combined by using a weights-based scheme in which each model contributes differently to the combined result. The weights are found by solving linear regression problems. Experiments conducted on two popular medical datasets namely CAMUS and Kvasir-SEG show that the proposed method achieves better results concerning two performance metrics (Dice Coefficient and Hausdorff distance) compared to some well-known benchmark algorithms.

研究动机与目标

• 解决在小规模医学图像数据集上训练的深度学习模型存在的过拟合和泛化能力有限的问题。
• 通过集成策略有效结合多个深度学习模型，以提升分割性能。
• 开发一种计算高效的多模型预测融合方法，利用学习到的权重进行融合。
• 通过分层集成架构，利用各模型的互补优势，弥补单个模型的弱点。
• 使用标准分割指标在基准医学图像数据集上验证该方法。

提出的方法

• 第一层由多个具有不同主干网络（如ResNet34、VGG16）的深度学习模型（如UNet、LinkNet、FPN）组成，这些模型在原始图像数据上进行训练。
• 将第一层模型的像素级预测输出与原始输入图像拼接，形成第二层模型的增强训练数据。
• 第二层模型在该增强数据上进行训练，利用第一层模型的综合信息来优化分割预测结果。
• 最终预测通过加权方案组合，其中类别特定权重通过线性回归学习，以最小化验证数据上的误差。
• 采用最小二乘回归方法对每种类别进行权重优化，以平衡各模型的贡献。
• 在两个公开数据集CAMUS和Kvasir-SEG上，使用Dice系数和Hausdorff距离评估该框架。

实验结果

研究问题

• RQ1 与单个模型相比，两层深度学习模型集成是否能提升在小规模医学图像数据集上的分割准确率？
• RQ2 将第一层模型的预测结果作为第二层的增强输入，是否能提升特征表示能力并增强模型泛化性能？
• RQ3 可学习的基于权重的融合策略是否能优于简单的平均或投票方法，以组合模型输出？
• RQ4 在超声心动图中，对于左心房和心肌等复杂解剖结构，该集成方法表现如何？
• RQ5 在边界检测误差方面，该集成方法在Hausdorff距离指标上的改善程度如何？

主要发现

• 在CAMUS数据集的心肌类别上，所提集成方法的Dice系数达到0.96，优于第二好的方法OLE（0.959）。
• 在左心房类别上，集成方法的Dice分数达到0.907，超过OLE（0.900）及其他模型，包括UNet-ResNet34（0.900）。
• 在收缩末期相，集成方法在左心室上的Hausdorff距离为4.4，优于第二好的模型LinkNet-ResNet34（4.7）。
• 在舒张末期相，集成方法在左心房上的Hausdorff距离为4.8，优于OLE（5.3）及其他模型。
• 该集成方法在复杂区域（如左心房右下角）的分割中优于单个模型预测，尤其在这些区域中，基础模型表现失败。
• 使用VGG16主干网络的模型表现较差，Dice分数低于0.3，凸显了主干网络选择及集成融合策略的重要性。