Main Content

使用深度学习进行三维脑肿瘤分割

此示例说明如何基于三维医学图像执行脑肿瘤的语义分割。

语义分割中,图像中的每个像素或三维体的每个体素都被标注为某一类。此示例说明如何在磁共振成像 (MRI) 扫描中使用三维 U-Net 深度学习网络执行脑肿瘤二元语义分割。U-Net 是一种快速、高效、简单的网络,在语义分割领域 [1] 非常流行。

医学图像分割面临的挑战之一是存储和处理三维体所需的内存量。由于 GPU 资源的限制,基于整个输入数据体训练网络并执行分割是不切实际的。此示例通过将图像分成更小的补片(或块)以训练和分割来解决问题。

医学图像分割面临的另一挑战是当使用常规的交叉熵损失时,数据中的类不平衡会妨碍训练。此示例通过使用加权多类 Dice 损失函数 [4] 解决此问题。对类进行加权有助于抵消较大区域对 Dice 分数的影响,使网络更容易学习如何分割较小区域。

此示例说明如何使用预训练的三维 U-Net 架构来执行脑肿瘤分割,以及如何使用一组测试图像来评估网络性能。您可以选择基于 BraTS 数据集 [2] 训练三维 U-Net。

使用预训练的三维 U-Net 执行脑肿瘤分割

下载预训练的三维 U-Net

将预训练的三维 U-Net 下载到名为 net 的变量中。

dataDir = fullfile(tempdir,"BraTS");
if ~exist(dataDir,'dir')
    mkdir(dataDir);
end
trained3DUnetURL = "https://www.mathworks.com/supportfiles/"+ ...
    "vision/data/brainTumor3DUNetValid.mat";
downloadTrainedNetwork(trained3DUnetURL,dataDir);
load(dataDir+filesep+"brainTumor3DUNetValid.mat");

下载 BraTS 样本数据

使用 downloadBraTSSampleTestData 辅助函数 [3],从 BraTS 数据集中下载五个样本测试数据体及其对应的标签。此辅助函数作为支持文件包含在本示例中。样本数据使您能够对测试数据执行分割,而无需下载完整数据集。

downloadBraTSSampleTestData(dataDir);

加载其中一个数据体样本及其像素标签真实值。

testDir = dataDir+filesep+"sampleBraTSTestSetValid";
data = load(fullfile(testDir,"imagesTest","BraTS446.mat"));
labels = load(fullfile(testDir,"labelsTest","BraTS446.mat"));
volTest = data.cropVol;
volTestLabels = labels.cropLabel;

执行语义分割

该示例使用重叠分块策略来处理大型数据体。重叠分块策略选择重叠的块,通过使用 semanticseg (Computer Vision Toolbox) 函数预测每个块的标签,然后将这些块重新组合成一个完整的分割测试数据体。该策略有助于在内存资源有限的 GPU 上实现高效处理。该策略还通过使用神经网络中卷积的有效部分来减少边界伪影 [5]。

通过将数据体数据存储为 blockedImage (Image Processing Toolbox) 对象并使用 apply (Image Processing Toolbox) 函数处理块来实现重叠分块策略。

为在上一节中下载的样本数据体创建一个 blockedImage 对象。

bim = blockedImage(volTest);

apply 函数对 blockedImage 中的每个块都执行自定义函数。将 semanticsegBlock 定义为对每个块都执行的函数。

semanticsegBlock = @(bstruct)semanticseg(bstruct.Data,net);

指定块大小作为网络输出大小。要创建重叠块,请指定非零的边界大小。此示例使用边界大小,使块加边界的大小等于网络输入大小。

networkInputSize = net.Layers(1).InputSize;
networkOutputSize = net.Layers(end).OutputSize;
blockSize = [networkOutputSize(1:3) networkInputSize(end)];
borderSize = (networkInputSize(1:3) - blockSize(1:3))/2;

在不完整块填充设置为 true 的情况下,使用 blockedImage apply 执行语义分割。默认填充方法 "replicate" 就很适用,因为数据体数据包含多个模态。批量大小指定为 1,以防止内存资源有限的 GPU 上出现内存不足错误。但是,如果您的 GPU 有足够的内存,则您可以通过增大块的大小来提高处理速度。

batchSize = 1;
results = apply(bim, ...
    semanticsegBlock, ...
    BlockSize=blockSize, ...
    BorderSize=borderSize,...
    PadPartialBlocks=true, ...
    BatchSize=batchSize);
predictedLabels = results.Source;

使用蒙太奇显示真实值和预测标签沿深度方向的中心切片。

zID = size(volTest,3)/2;
zSliceGT = labeloverlay(volTest(:,:,zID),volTestLabels(:,:,zID));
zSlicePred = labeloverlay(volTest(:,:,zID),predictedLabels(:,:,zID));

figure
montage({zSliceGT,zSlicePred},Size=[1 2],BorderSize=5) 
title("Labeled Ground Truth (Left) vs. Network Prediction (Right)")

Figure contains an axes object. The axes object with title Labeled Ground Truth (Left) vs. Network Prediction (Right) contains an object of type image.

以下图像顺序显示了该三维体的各切片。标注的真实值在左侧,网络预测在右侧。

训练三维 U-Net

该示例的此部分显示如何训练三维 U-Net。如果您不想下载训练数据集或训练网络,则可以跳到此示例的评估网络性能部分。

下载 BraTS 数据集

此示例使用 BraTS 数据集 [2]。BraTS 数据集包含脑肿瘤(即胶质瘤,最常见的原发性恶性脑肿瘤)的 MRI 扫描。该数据文件的大小约为 7 GB。

要下载 BraTS 数据,请访问 Medical Segmentation Decathlon 网站,然后点击“Download Data”链接。下载 "Task01_BrainTumour.tar" 文件 [3]。将该 TAR 文件解压缩到由 imageDir 变量指定的目录中。成功解压缩后,imageDir 将包含一个名为 Task01_BrainTumour 的目录,该目录有三个子目录:imagesTrimagesTslabelsTr

该数据集包含 750 个四维体,每个四维体代表一组三维图像。每个四维体的大小为 240×240×155×4,其中前三个维度对应于三维体图像的高度、宽度和深度。第四个维度对应于不同的扫描形态。该数据集分为 484 个带体素标签的训练数据体和 266 个测试数据体。测试数据体没有标签,因此此示例不使用测试数据。取而代之,示例将 484 个训练数据体拆分成三个独立的数据集,分别用于训练、验证和测试。

预处理训练和验证数据

为了更高效地训练三维 U-Net 网络,使用辅助函数 preprocessBraTSDataset 预处理 MRI 数据。此函数作为支持文件包含在本示例中。该辅助函数执行以下操作:

  • 将数据裁剪到主要包含大脑和肿瘤的区域。裁剪数据可以减小数据的大小,同时保留每个 MRI 数据体的最关键部分及其对应标签。

  • 通过减去均值并除以裁剪后的大脑区域的标准差,独立地对每个数据体的每个形态进行归一化。

  • 将 484 个训练数据体拆分成 400 个训练集、29 个验证集和 55 个测试集。

数据预处理可能需要大约 30 分钟才能完成。

sourceDataLoc = dataDir+filesep+"Task01_BrainTumour";
preprocessDataLoc = dataDir+filesep+"preprocessedDataset";
preprocessBraTSDataset(preprocessDataLoc,sourceDataLoc);

为训练和验证创建随机补片提取数据存储

创建一个 imageDatastore 来存储三维图像数据。由于 MAT 文件格式是非标准图像格式,您必须使用 MAT 文件读取器来读取图像数据。您可以使用 MAT 文件读取器辅助函数 matRead。此函数作为支持文件包含在本示例中。

volLoc = fullfile(preprocessDataLoc,"imagesTr");
volds = imageDatastore(volLoc,FileExtensions=".mat",ReadFcn=@matRead);

创建一个 pixelLabelDatastore (Computer Vision Toolbox) 来存储标签。

lblLoc = fullfile(preprocessDataLoc,"labelsTr");
classNames = ["background","tumor"];
pixelLabelID = [0 1];
pxds = pixelLabelDatastore(lblLoc,classNames,pixelLabelID, ...
    FileExtensions=".mat",ReadFcn=@matRead);

创建一个 randomPatchExtractionDatastore (Image Processing Toolbox),它从真实值图像提取随机补片和对应的像素标签数据。指定 132×132×132 体素的补片大小。指定 "PatchesPerImage" 以在训练期间从每对数据体和标签中提取 16 个随机定位的补片。指定小批量大小为 8。

patchSize = [132 132 132];
patchPerImage = 16;
miniBatchSize = 8;
patchds = randomPatchExtractionDatastore(volds,pxds,patchSize, ...
    PatchesPerImage=patchPerImage);
patchds.MiniBatchSize = miniBatchSize;

创建一个 randomPatchExtractionDatastore,它从验证图像和像素标签数据中提取补片。您可以使用验证数据来评估网络随着时间的推移是在持续学习、欠拟合还是过拟合。

volLocVal = fullfile(preprocessDataLoc,"imagesVal");
voldsVal = imageDatastore(volLocVal,FileExtensions=".mat", ...
    ReadFcn=@matRead);

lblLocVal = fullfile(preprocessDataLoc,"labelsVal");
pxdsVal = pixelLabelDatastore(lblLocVal,classNames,pixelLabelID, ...
    FileExtensions=".mat",ReadFcn=@matRead);

dsVal = randomPatchExtractionDatastore(voldsVal,pxdsVal,patchSize, ...
    PatchesPerImage=patchPerImage);
dsVal.MiniBatchSize = miniBatchSize;

设置三维 U-Net 层

此示例使用三维 U-Net 网络 [1]。U-Net 中的初始卷积层序列与最大池化层交叠,从而逐步降低输入图像的分辨率。这些层后跟一系列使用上采样算子散布的卷积层,从而会连续增加输入图像的分辨率。在每个 ReLU 层之前引入一个批量归一化层。U-Net 的名称源于网络可以绘制成形似字母 U 的对称形状。

使用 unetLayers (Computer Vision Toolbox) 函数创建一个默认的三维 U-Net 网络。指定二类分割。还要指定有效的卷积填充,以避免在使用重叠分块策略预测测试数据体时出现边界伪影。

numChannels = 4;
inputPatchSize = [patchSize numChannels];
numClasses = 2;
[lgraph,outPatchSize] = unet3dLayers(inputPatchSize, ...
    numClasses,ConvolutionPadding="valid");

通过使用 transform 函数和辅助函数 augmentAndCrop3dPatch 指定的自定义预处理操作来增强训练数据和验证数据。此函数作为支持文件包含在本示例中。augmentAndCrop3dPatch 函数执行以下操作:

  1. 随机旋转和翻转训练数据,使训练更加稳健。该函数不旋转或翻转验证数据。

  2. 将响应补片裁剪为网络的输出大小 44×44×44 体素。

dsTrain = transform(patchds, ...
    @(patchIn)augmentAndCrop3dPatch(patchIn,outPatchSize,"Training"));
dsVal = transform(dsVal, ...
    @(patchIn)augmentAndCrop3dPatch(patchIn,outPatchSize,"Validation"));

为了更好地分割较小的肿瘤区域并减少较大背景区域的影响,此示例使用 dicePixelClassificationLayer (Computer Vision Toolbox)。将像素分类层替换为 Dice 像素分类层。

outputLayer = dicePixelClassificationLayer(Name="Output");
lgraph = replaceLayer(lgraph,"Segmentation-Layer",outputLayer);

数据已在此示例的预处理训练和验证数据部分进行归一化。image3dInputLayer 中没有必要进行数据归一化,因此将输入层替换为没有数据归一化的输入层。

inputLayer = image3dInputLayer(inputPatchSize, ...
    Normalization="none",Name="ImageInputLayer");
lgraph = replaceLayer(lgraph,"ImageInputLayer",inputLayer);

您也可以使用深度网络设计器修改三维 U-Net 网络。

deepNetworkDesigner(lgraph)

指定训练选项

使用 adam 优化求解器来训练网络。使用 trainingOptions 函数指定超参数设置。初始学习率设置为 5e-4,并在训练期间逐渐降低。您可以根据您的 GPU 内存情况尝试调整 MiniBatchSize 属性。为了最大限度地利用 GPU 内存,最好使用大输入补片而不是大批量大小。请注意,批量归一化层对于较小的 MiniBatchSize 值不是很有效。根据 MiniBatchSize 调整初始学习率。

options = trainingOptions("adam", ...
    MaxEpochs=50, ...
    InitialLearnRate=5e-4, ...
    LearnRateSchedule="piecewise", ...
    LearnRateDropPeriod=5, ...
    LearnRateDropFactor=0.95, ...
    ValidationData=dsVal, ...
    ValidationFrequency=400, ...
    Plots="training-progress", ...
    Verbose=false, ...
    MiniBatchSize=miniBatchSize);

训练网络

默认情况下,该示例使用下载的预训练三维 U-Net 网络。借助预训练网络,您无需等待训练完成,即可执行语义分割和评估分割结果。

要训练网络,请将以下代码中的 doTraining 变量设置为 true。使用 trainNetwork 函数训练网络。

在 GPU 上(如果有)进行训练。使用 GPU 需要 Parallel Computing Toolbox™ 和支持 CUDA® 的 NVIDIA® GPU。有关详细信息,请参阅GPU Computing Requirements (Parallel Computing Toolbox)。在包含 4 个 NVIDIA™ Titan Xp GPU 的多 GPU 系统上进行训练需要大约 30 个小时,根据您的 GPU 硬件情况,训练时间可能会更长。

doTraining = false;
if doTraining
    [net,info] = trainNetwork(dsTrain,lgraph,options);
    modelDateTime = string(datetime("now",Format="yyyy-MM-dd-HH-mm-ss"));
    save("trained3DUNet-"+modelDateTime+".mat","net");
end

评估网络性能

选择包含真实值数据体和测试标签的测试数据源。如果您在以下代码中将 useFullTestSet 变量保留为 false,则该示例使用五个样本数据体进行测试。如果将 useFullTestSet 变量设置为 true,则该示例使用从完整数据集选择的 55 个测试图像。

useFullTestSet = false;
if useFullTestSet
    volLocTest = fullfile(preprocessDataLoc,"imagesTest");
    lblLocTest = fullfile(preprocessDataLoc,"labelsTest");
else
    volLocTest = fullfile(testDir,"imagesTest");
    lblLocTest = fullfile(testDir,"labelsTest");
end

voldsTest 变量存储真实值测试图像。pxdsTest 变量存储真实值标签。

voldsTest = imageDatastore(volLocTest,FileExtensions=".mat", ...
    ReadFcn=@matRead);
pxdsTest = pixelLabelDatastore(lblLocTest,classNames,pixelLabelID, ...
    FileExtensions=".mat",ReadFcn=@matRead);

对于每个测试数据体,使用 apply (Image Processing Toolbox) 函数处理每个块。apply 函数执行由辅助函数 calculateBlockMetrics 指定的运算,该辅助函数在此示例的末尾定义。calculateBlockMetrics 函数执行每个块的语义分割,并计算预测标签和真实值标签之间的混淆矩阵。

imageIdx = 1;
datasetConfMat = table;
while hasdata(voldsTest)

    % Read volume and label data
    vol = read(voldsTest);
    volLabels = read(pxdsTest);

    % Create blockedImage for volume and label data
    testVolume = blockedImage(vol);
    testLabels = blockedImage(volLabels{1});

    % Calculate block metrics
    blockConfMatOneImage = apply(testVolume, ...
        @(block,labeledBlock) ...
            calculateBlockMetrics(block,labeledBlock,net), ...
        ExtraImages=testLabels, ...
        PadPartialBlocks=true, ...
        BlockSize=blockSize, ...
        BorderSize=borderSize, ...
        UseParallel=false);

    % Read all the block results of an image and update the image number
    blockConfMatOneImageDS = blockedImageDatastore(blockConfMatOneImage);
    blockConfMat = readall(blockConfMatOneImageDS);
    blockConfMat = struct2table([blockConfMat{:}]);
    blockConfMat.ImageNumber = imageIdx.*ones(height(blockConfMat),1);
    datasetConfMat = [datasetConfMat;blockConfMat];

    imageIdx = imageIdx + 1;
end

使用 evaluateSemanticSegmentation (Computer Vision Toolbox) 函数评估分割的数据集指标和块指标。

[metrics,blockMetrics] = evaluateSemanticSegmentation( ...
    datasetConfMat,classNames,Metrics="all");
Evaluating semantic segmentation results
----------------------------------------
* Selected metrics: global accuracy, class accuracy, IoU, weighted IoU.
* Processed 5 images.
* Finalizing... Done.
* Data set metrics:

    GlobalAccuracy    MeanAccuracy    MeanIoU    WeightedIoU
    ______________    ____________    _______    ___________

       0.99902          0.97955       0.95978      0.99808  

显示为每个图像计算的 Jaccard 分数。

metrics.ImageMetrics.MeanIoU
ans = 5×1

    0.9613
    0.9570
    0.9551
    0.9656
    0.9594

支持函数

calculateBlockMetrics 辅助函数执行块的语义分割,并计算预测标签和真实值标签之间的混淆矩阵。该函数返回一个结构体,其字段包含关于该块的混淆矩阵和元数据。您可以将该结构体与 evaluateSemanticSegmentation 函数结合使用来计算指标和聚合基于块的结果。

function blockMetrics = calculateBlockMetrics(bstruct,gtBlockLabels,net)

% Segment block
predBlockLabels = semanticseg(bstruct.Data,net);

% Trim away border region from gtBlockLabels 
blockStart = bstruct.BorderSize + 1;
blockEnd = blockStart + bstruct.BlockSize - 1;
gtBlockLabels = gtBlockLabels( ...
    blockStart(1):blockEnd(1), ...
    blockStart(2):blockEnd(2), ...
    blockStart(3):blockEnd(3));

% Evaluate segmentation results against ground truth
confusionMat = segmentationConfusionMatrix(predBlockLabels,gtBlockLabels);

% blockMetrics is a struct with confusion matrices, image number,
% and block information. 
blockMetrics.ConfusionMatrix = confusionMat;
blockMetrics.ImageNumber = bstruct.ImageNumber;
blockInfo.Start = bstruct.Start;
blockInfo.End = bstruct.End;
blockMetrics.BlockInfo = blockInfo;

end

参考资料

[1] Çiçek, Ö., A. Abdulkadir, S. S. Lienkamp, T. Brox, and O. Ronneberger."3D U-Net:Learning Dense Volumetric Segmentation from Sparse Annotation."In Proceedings of the International Conference on Medical Image Computing and Computer-Assisted Intervention - MICCAI 2016.Athens, Greece, Oct. 2016, pp. 424-432.

[2] Isensee, F., P. Kickingereder, W. Wick, M. Bendszus, and K. H. Maier-Hein."Brain Tumor Segmentation and Radiomics Survival Prediction:Contribution to the BRATS 2017 Challenge."In Proceedings of BrainLes:International MICCAI Brainlesion Workshop.Quebec City, Canada, Sept. 2017, pp. 287-297.

[3] "Brain Tumours".Medical Segmentation Decathlon. http://medicaldecathlon.com/

BraTS 数据集由 Medical Segmentation Decathlon 在 CC-BY-SA 4.0 许可证下提供。所有保证和陈述均有相应的免责声明;有关详细信息,请参阅许可文档。MathWorks® 已修改此示例的下载 BraTS 样本数据部分中链接的数据集。修改后的样本数据集已裁剪为主要包含大脑和肿瘤的区域,并且通过减去均值并除以裁剪后的大脑区域的标准差,独立地对每个通道进行归一化。

[4] Sudre, C. H., W. Li, T. Vercauteren, S. Ourselin, and M. J. Cardoso."Generalised Dice Overlap as a Deep Learning Loss Function for Highly Unbalanced Segmentations."Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support:Third International Workshop.Quebec City, Canada, Sept. 2017, pp. 240-248.

[5] Ronneberger, O., P. Fischer, and T. Brox."U-Net:Convolutional Networks for Biomedical Image Segmentation."In Proceedings of the International Conference on Medical Image Computing and Computer-Assisted Intervention - MICCAI 2015.Munich, Germany, Oct. 2015, pp. 234-241.Available at arXiv:1505.04597.

另请参阅

(Image Processing Toolbox) | | | | (Computer Vision Toolbox) | | (Computer Vision Toolbox) | (Computer Vision Toolbox)

相关主题