Main Content

本页的翻译已过时。点击此处可查看最新英文版本。

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

此示例说明如何基于三维医学图像训练三维 U-Net 神经网络,并执行脑肿瘤的语义分割。此示例说明如何训练三维 U-Net 网络,并提供一个预训练网络。强烈建议使用具有 3.0 或更高计算能力的、支持 CUDA 的 NVIDIA™ GPU 来进行三维语义分割(需要 Parallel Computing Toolbox™)。

简介

语义分割中,图像中的每个像素或三维体的每个体素都被标注为某一类。此示例说明如何在磁共振成像 (MRI) 扫描中使用深度学习方法执行脑肿瘤二元语义分割。在该二值分割中,将每个像素标注为肿瘤或背景。

此示例使用三维 U-Net 架构 [1] 执行脑肿瘤分割。U-Net 是一种快速、高效、简单的网络,在语义分割领域非常流行。

医学图像分割面临的挑战之一是存储和处理三维体所需的内存量。由于 GPU 资源的限制,基于整个输入图像体训练网络是不切实际的。此示例通过基于图像补片训练网络来解决此问题。该示例使用重叠分块策略,将测试补片缝合为完整的经过分割的测试图像体。该示例使用神经网络中卷积的有效部分来避免边界伪影 [5]。

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

下载训练、验证和测试数据

此示例使用 BraTS 数据集 [2]。BraTS 数据集包含脑肿瘤(即胶质瘤,最常见的原发性恶性脑肿瘤)的 MRI 扫描。该数据文件的大小约为 7 GB。如果您不想下载 BraTS 数据集,请直接转至本示例中的下载预训练网络和样本测试集部分。

创建一个目录以存储 BraTS 数据集。

imageDir = fullfile(tempdir,'BraTS');
if ~exist(imageDir,'dir')
    mkdir(imageDir);
end

要下载 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 = [imageDir filesep 'Task01_BrainTumour'];
preprocessDataLoc = fullfile(tempdir,'BraTS','preprocessedDataset');
preprocessBraTSdataset(preprocessDataLoc,sourceDataLoc);

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

使用随机补片提取数据存储将训练数据馈送到网络,并验证训练进度。此数据存储从真实值图像提取随机补片和对应的像素标签数据。补片是一种常用方法,用于在使用任意大的图像体训练时防止内存耗尽。

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

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

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

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

预览一个带标签的图像体。使用 labelvolshow (Image Processing Toolbox) 函数显示标注图像体。通过将背景标签的可见性 (1) 设置为 0,使背景完全透明。

volume = preview(volds);
label = preview(pxds);

viewPnl = uipanel(figure,'Title','Labeled Training Volume');
hPred = labelvolshow(label,volume(:,:,:,1),'Parent',viewPnl, ...
    'LabelColor',[0 0 0;1 0 0]);
hPred.LabelVisibility(1) = 0;

创建一个包含训练图像和像素标签数据的 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',volReader);

lblLocVal = fullfile(preprocessDataLoc,'labelsVal');
pxdsVal = pixelLabelDatastore(lblLocVal,classNames,pixelLabelID, ...
    'FileExtensions','.mat','ReadFcn',volReader);

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

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

augmentAndCrop3dPatch 函数执行以下操作:

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

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

dataSource = 'Training';
dsTrain = transform(patchds,@(patchIn)augmentAndCrop3dPatch(patchIn,dataSource));

dataSource = 'Validation';
dsVal = transform(dsVal,@(patchIn)augmentAndCrop3dPatch(patchIn,dataSource));

设置三维 U-Net 层

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

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

inputPatchSize = [132 132 132 4];
numClasses = 2;
[lgraph,outPatchSize] = unet3dLayers(inputPatchSize,numClasses,'ConvolutionPadding','valid');

为了更好地分割较小的肿瘤区域并减少较大背景区域的影响,此示例使用 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);

或者,您可以使用 Deep Learning Toolbox™ 中的深度网络设计器来修改三维 U-Net 网络。

绘制更新后的三维 U-Net 网络图。

analyzeNetwork(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 和 BraTS 数据集 [3] 中的五个样本测试图像体及其对应的标签。预训练的模型和样本数据使您能够对测试数据执行分割,而无需下载完整的数据集或等待网络完成训练。

trained3DUnet_url = 'https://www.mathworks.com/supportfiles/vision/data/brainTumor3DUNetValid.mat';
sampleData_url = 'https://www.mathworks.com/supportfiles/vision/data/sampleBraTSTestSetValid.tar.gz';

imageDir = fullfile(tempdir,'BraTS');
if ~exist(imageDir,'dir')
    mkdir(imageDir);
end

downloadTrained3DUnetSampleData(trained3DUnet_url,sampleData_url,imageDir);
Downloading pretrained 3-D U-Net for BraTS data set.
This will take several minutes to download...
Done.

Downloading sample BraTS test dataset.
This will take several minutes to download and unzip...
Done.

训练网络

在配置训练选项和数据源后,使用 trainNetwork 函数训练三维 U-Net 网络。要训练网络,请将以下代码中的 doTraining 变量设置为 true。强烈建议使用具有 3.0 或更高计算能力的、支持 CUDA 的 NVIDIA™ GPU 来进行训练。

如果将以下代码中的 doTraining 变量保留为 false,则该示例返回预训练的三维 U-Net 网络。

注意:在包含 4 个 NVIDIA™ Titan Xp GPU 的多 GPU 系统上进行训练需要大约 30 个小时,根据您的 GPU 硬件情况,训练时间可能会更长。

doTraining = false;
if doTraining
    modelDateTime = datestr(now,'dd-mmm-yyyy-HH-MM-SS');
    [net,info] = trainNetwork(dsTrain,lgraph,options);
    save(['trained3DUNetValid-' modelDateTime '-Epoch-' num2str(options.MaxEpochs) '.mat'],'net');
else
    inputPatchSize = [132 132 132 4];
    outPatchSize = [44 44 44 2];
    load(fullfile(imageDir,'trained3DUNet','brainTumor3DUNetValid.mat'));
end

您现在可以使用 U-Net 对脑肿瘤进行语义分割。

对测试数据执行分割

强烈建议使用 GPU 来执行图像体的语义分割(需要 Parallel Computing Toolbox™)。

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

useFullTestSet = false;
if useFullTestSet
    volLocTest = fullfile(preprocessDataLoc,'imagesTest');
    lblLocTest = fullfile(preprocessDataLoc,'labelsTest');
else
    volLocTest = fullfile(imageDir,'sampleBraTSTestSetValid','imagesTest');
    lblLocTest = fullfile(imageDir,'sampleBraTSTestSetValid','labelsTest');
    classNames = ["background","tumor"];
    pixelLabelID = [0 1];
end

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

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

使用重叠分块策略来预测每个测试图像体的标签。对每个测试图像体都进行填充,以使输入大小是网络输出大小的倍数,并补偿有效卷积的效应。重叠分块算法选择重叠补片,使用 semanticseg (Computer Vision Toolbox) 函数预测每个补片的标签,然后重新组合补片。

id = 1;
while hasdata(voldsTest)
    disp(['Processing test volume ' num2str(id)]);
    
    tempGroundTruth = read(pxdsTest);
    groundTruthLabels{id} = tempGroundTruth{1};
    vol{id} = read(voldsTest);
    
    % Use reflection padding for the test image. 
    % Avoid padding of different modalities.
    volSize = size(vol{id},(1:3));
    padSizePre  = (inputPatchSize(1:3)-outPatchSize(1:3))/2;
    padSizePost = (inputPatchSize(1:3)-outPatchSize(1:3))/2 + (outPatchSize(1:3)-mod(volSize,outPatchSize(1:3)));
    volPaddedPre = padarray(vol{id},padSizePre,'symmetric','pre');
    volPadded = padarray(volPaddedPre,padSizePost,'symmetric','post');
    [heightPad,widthPad,depthPad,~] = size(volPadded);
    [height,width,depth,~] = size(vol{id});
    
    tempSeg = categorical(zeros([height,width,depth],'uint8'),[0;1],classNames);
    
    % Overlap-tile strategy for segmentation of volumes.
    for k = 1:outPatchSize(3):depthPad-inputPatchSize(3)+1
        for j = 1:outPatchSize(2):widthPad-inputPatchSize(2)+1
            for i = 1:outPatchSize(1):heightPad-inputPatchSize(1)+1
                patch = volPadded( i:i+inputPatchSize(1)-1,...
                    j:j+inputPatchSize(2)-1,...
                    k:k+inputPatchSize(3)-1,:);
                patchSeg = semanticseg(patch,net);
                tempSeg(i:i+outPatchSize(1)-1, ...
                    j:j+outPatchSize(2)-1, ...
                    k:k+outPatchSize(3)-1) = patchSeg;
            end
        end
    end
    
    % Crop out the extra padded region.
    tempSeg = tempSeg(1:height,1:width,1:depth);

    % Save the predicted volume result.
    predictedLabels{id} = tempSeg;
    id=id+1;
end
Processing test volume 1
Processing test volume 2
Processing test volume 3
Processing test volume 4
Processing test volume 5

将真实值与网络预测进行比较

选择测试图像之一来评估语义分割的准确度。从四维体数据中提取第一个形态,并将此三维体存储在变量 vol3d 中。

volId = 1;
vol3d = vol{volId}(:,:,:,1);

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

zID = size(vol3d,3)/2;
zSliceGT = labeloverlay(vol3d(:,:,zID),groundTruthLabels{volId}(:,:,zID));
zSlicePred = labeloverlay(vol3d(:,:,zID),predictedLabels{volId}(:,:,zID));

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

使用 labelvolshow (Image Processing Toolbox) 函数显示标注真实值的图像体。通过将背景标签的可见性 (1) 设置为 0,使背景完全透明。由于肿瘤在脑组织内部,因此要使一些大脑体素透明,以便肿瘤可见。要使一些大脑体素透明,请将图像体阈值指定为 [0, 1] 范围内的一个数字。归一化的图像体强度低于此阈值时即完全透明。此示例将图像体阈值设置为小于 1,以使一些大脑像素保持可见,从而给出肿瘤在大脑中所处空间位置的上下文。

viewPnlTruth = uipanel(figure,'Title','Ground-Truth Labeled Volume');
hTruth = labelvolshow(groundTruthLabels{volId},vol3d,'Parent',viewPnlTruth, ...
    'LabelColor',[0 0 0;1 0 0],'VolumeThreshold',0.68);
hTruth.LabelVisibility(1) = 0;

对于同一图像体,显示预测的标签。

viewPnlPred = uipanel(figure,'Title','Predicted Labeled Volume');
hPred = labelvolshow(predictedLabels{volId},vol3d,'Parent',viewPnlPred, ...
    'LabelColor',[0 0 0;1 0 0],'VolumeThreshold',0.68);

hPred.LabelVisibility(1) = 0;

以下图像展示了逐层显示一个图像体切片的过程。标注的真实值在左侧,网络预测在右侧。

量化分割准确度

使用 dice (Image Processing Toolbox) 函数测量分割准确度。此函数计算预测分割和真实值分割之间的 Dice 相似性系数。

diceResult = zeros(length(voldsTest.Files),2);

for j = 1:length(vol)
    diceResult(j,:) = dice(groundTruthLabels{j},predictedLabels{j});
end

计算整个测试图像体集合的平均 Dice 分数。

meanDiceBackground = mean(diceResult(:,1));
disp(['Average Dice score of background across ',num2str(j), ...
    ' test volumes = ',num2str(meanDiceBackground)])
Average Dice score of background across 5 test volumes = 0.9993
meanDiceTumor = mean(diceResult(:,2));
disp(['Average Dice score of tumor across ',num2str(j), ...
    ' test volumes = ',num2str(meanDiceTumor)])
Average Dice score of tumor across 5 test volumes = 0.9585

下图显示一个 boxplot (Statistics and Machine Learning Toolbox),它可视化包含五个样本测试图像体的集合中 Dice 分数的统计量。图中的红线显示类的 Dice 值中位数。蓝框的上界和下界分别表示第 25 个和第 75 个百分位数。黑色须线会延伸到不是离群值的最远端数据点。

如果您有 Statistics and Machine Learning Toolbox™,则可以使用 boxplot 函数来可视化所有测试图像体的 Dice 分数的统计量。要创建一个 boxplot,请将以下代码中的 createBoxplot 变量设置为 true

createBoxplot = false;
if createBoxplot
    figure
    boxplot(diceResult)
    title('Test Set Dice Accuracy')
    xticklabels(classNames)
    ylabel('Dice Coefficient')
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® 已修改此示例的下载预训练网络和样本测试集部分中链接的数据集。修改后的样本数据集已裁剪为主要包含大脑和肿瘤的区域,并且通过减去均值并除以裁剪后的大脑区域的标准差,独立地对每个通道进行归一化。

[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.

另请参阅

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

相关主题