Main Content

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

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

此示例说明如何基于三维医学图像训练三维 U-Net 神经网络,并执行脑肿瘤的语义分割。

语义分割中,图像中的每个像素或三维体的每个体素都被标注为某一类。此示例说明如何在磁共振成像 (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);

训练网络

默认情况下,该示例加载一个预训练的三维 U-Net 网络。借助预训练网络,您无需等待训练完成,即可运行整个示例。

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

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

doTraining = false;
if doTraining
    modelDateTime = string(datetime('now','Format',"yyyy-MM-dd-HH-mm-ss"));
    [net,info] = trainNetwork(dsTrain,lgraph,options);
    save(strcat("trained3DUNet-",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

对测试数据执行分割

强烈建议使用 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.

另请参阅

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

相关主题