本页对应的英文页面已更新,但尚未翻译。 若要查看最新内容,请点击此处访问英文页面。

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

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

此示例说明如何训练三维 U-Net 网络,并提供一个预训练网络。如果您选择自己训练网络,强烈建议使用计算能力为 3.0 或更高的支持 CUDA 的 NVIDIA™ GPU(需要 Parallel Computing Toolbox™)。

如果您不想下载整个训练数据集或自己训练网络,则可以加载预训练的三维 U-Net 和较小的测试数据集。请直接转至本示例中的下载预训练网络和样本测试集部分。

简介

语义分割中,图像中的每个像素或三维体的每个体素都被标记为某一类。

此示例说明如何在磁共振成像 (MRI) 扫描中使用深度学习方法对脑肿瘤进行语义分割。它还说明如何执行二元分割,将每个体素标记为肿瘤或背景。

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

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

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

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

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

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

要下载 BraTS 数据,请访问 Medical Segmentation Decathalon 网站,然后点击 "Download Data" 链接。下载 Task01_BrainTumour.tar 文件。将该 TAR 文件解压缩到由 imageDir 变量指定的目录中。

该数据集包含 750 个四维体,每个四维体代表一组三维图像。每个四维体的大小为 240×240×155×4,其中前三个维度对应于三维体图像的高度、宽度和深度。第四个维度对应于不同的扫描形态。每个图像体都有一个对应的像素标签。该数据集分为 484 个训练四维体和 286 个测试四维体。

预处理训练和验证数据

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

该辅助函数执行以下操作:

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

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

  • 将数据集分成训练集、验证集和测试集。

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

preprocessDataLoc = fullfile(tempdir,'BraTS','preprocessedDataset');
preprocessBraTSdataset(preprocessDataLoc,imageDir);

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

使用随机补片提取数据存储将训练数据馈送到网络,并验证训练进度。此数据存储从真实值图像提取随机补片和对应的像素标签数据。

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

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

创建一个 pixelLabelDatastore 来存储标签。

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

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

volume = preview(volds);
label = preview(pxds);
h = labelvolshow(label,volume(:,:,:,1));
h.LabelVisibility(1) = 0;

基于图像数据存储和像素标签数据存储创建一个 randomPatchExtractionDatastore。指定 64×64×64 体素的补片大小。指定 'PatchesPerImage' 以在训练期间从每对图像体和标签中提取 16 个随机定位的补片。指定小批量大小为 8。此步骤是必需的,这样 GPU 资源就可以高效地处理任意大的图像体而不会耗尽内存。

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

通过使用 transform 函数和辅助函数 augment3dPatch 指定的自定义预处理操作来增强训练数据。augment3dPatch 函数随机旋转和翻转训练数据,使训练更加稳健。此函数作为支持文件包含在本示例中。

augpatchds = transform(patchds,@augment3dPatch);

创建一个包含验证数据的 randomPatchExtrationDatastore。您可以使用验证数据来评估网络随着时间的推移是在持续学习、欠拟合还是过拟合。

volLocVal = fullfile(preprocessDataLoc,'imagesVal');
voldsVal = imageDatastore(volLocVal, ...
    'FileExtensions','.mat','ReadFcn',volReader);

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

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

设置三维 U-Net 层

此示例使用三维 U-Net 网络 [1] 的一个变体。U-Net 中的初始卷积层序列与最大池化层交叠,从而逐步降低输入图像的分辨率。这些层后跟一系列使用上采样算子散布的卷积层,从而会连续增加输入图像的分辨率。U-Net 的名称源于网络可以绘制成形似字母 U 的对称形状。此示例对 U-Net 进行了修改以在卷积中使用填零,从而使卷积的输入和输出具有相同的大小。

此示例使用 Deep Learning Toolbox™ 中的层定义三维 U-Net,这些层包括:

此示例还定义一个名为 dicePixelClassidication3dLayer 的自定义 Dice 损失层,以解决数据中的类不平衡问题。此层作为支持文件包含在本示例中。有关详细信息,请参阅定义使用 Dice 损失的自定义像素分类层

第一个层 imageInput3dLayer 对大小为 64×64×64 体素的图像补片进行操作。

inputSize = [64 64 64 4];
inputLayer = image3dInputLayer(inputSize,'Normalization','none','Name','input');

图像输入层后跟三维 U-Net 的收缩路径。收缩路径由三个编码器模块组成。每个编码器包含两个具有 3×3×3 过滤器的卷积层,使特征映射的数量翻倍,后跟一个用作非线性激活的 reLu 层。第一个卷积层还后跟一个批量归一化层。每个编码器以一个最大池化层结束,该层将图像分辨率在每个维度上减半。

为所有层指定唯一名称。编码器中各层的名称以 "en" 开头,后跟编码器模块的索引。例如,"en1" 表示第一个编码器模块。

numFiltersEncoder = [
    32 64; 
    64 128; 
    128 256];

layers = [inputLayer];
for module = 1:3
    modtag = num2str(module);
    encoderModule = [
        convolution3dLayer(3,numFiltersEncoder(module,1), ...
            'Padding','same','WeightsInitializer','normal', ...
            'Name',['en',modtag,'_conv1']);
        batchNormalizationLayer('Name',['en',modtag,'_bn']);
        reluLayer('Name',['en',modtag,'_relu1']);
        convolution3dLayer(3,numFiltersEncoder(module,2), ...
            'Padding','same','WeightsInitializer','normal', ...
            'Name',['en',modtag,'_conv2']);
        reluLayer('Name',['en',modtag,'_relu2']);
        maxPooling3dLayer(2,'Stride',2,'Padding','same', ...
            'Name',['en',modtag,'_maxpool']);
    ];
    
    layers = [layers; encoderModule];
end

创建三维 U-Net 的扩展路径。扩展路径由四个解码器模块组成。所有编码器都包含两个具有 3×3×3 过滤器的卷积层,使特征映射的数量减半,后跟一个用作非线性激活的 reLu 层。前三个解码器以一个转置卷积层结束,该卷积层以 2 为因子对图像上采样。最终解码器则包含一个卷积层,该卷积层将每个体素的特征向量映射到类。

为所有层指定唯一名称。解码器中各层的名称以 "de" 开头,后跟解码器模块的索引。例如,"de4" 表示第四个解码器模块。

numFiltersDecoder = [
    256 512; 
    256 256; 
    128 128; 
    64 64];

decoderModule4 = [
    convolution3dLayer(3,numFiltersDecoder(1,1), ...
        'Padding','same','WeightsInitializer','normal', ...
        'Name','de4_conv1');
    reluLayer('Name','de4_relu1');
    convolution3dLayer(3,numFiltersDecoder(1,2), ...
        'Padding','same','WeightsInitializer','normal', ...
        'Name','de4_conv2');
    reluLayer('Name','de4_relu2');
    transposedConv3dLayer(2,numFiltersDecoder(1,2),'Stride',2, ...
        'Name','de4_transconv');
];

decoderModule3 = [
    convolution3dLayer(3,numFiltersDecoder(2,1), ...
        'Padding','same','WeightsInitializer','normal', ....
        'Name','de3_conv1');       
    reluLayer('Name','de3_relu1');
    convolution3dLayer(3,numFiltersDecoder(2,2), ...
        'Padding','same','WeightsInitializer','normal', ...
        'Name','de3_conv2'); 
    reluLayer('Name','de3_relu2');
    transposedConv3dLayer(2,numFiltersDecoder(2,2),'Stride',2, ...
        'Name','de3_transconv');
];

decoderModule2 = [
    convolution3dLayer(3,numFiltersDecoder(3,1), ...
        'Padding','same','WeightsInitializer','normal', ...
        'Name','de2_conv1');
    reluLayer('Name','de2_relu1');
    convolution3dLayer(3,numFiltersDecoder(3,2), ...
        'Padding','same','WeightsInitializer','normal', ...
        'Name','de2_conv2');
    reluLayer('Name','de2_relu2');
    transposedConv3dLayer(2,numFiltersDecoder(3,2),'Stride',2, ...
        'Name','de2_transconv');
];

最终解码器包含一个卷积层,该卷积层将每个体素的特征向量映射到两类(肿瘤和背景)之一。自定义 Dice 像素分类层对损失函数进行加权,以增加小肿瘤区域对 Dice 分数的影响。

numLabels = 2;
decoderModuleFinal = [
    convolution3dLayer(3,numFiltersDecoder(4,1), ...
        'Padding','same','WeightsInitializer','normal', ...
        'Name','de1_conv1');
    reluLayer('Name','de1_relu1');
    convolution3dLayer(3,numFiltersDecoder(4,2), ...
        'Padding','same','WeightsInitializer','normal', ...
        'Name','de1_conv2');
    reluLayer('Name','de1_relu2');
    convolution3dLayer(1,numLabels,'Name','convLast');
    softmaxLayer('Name','softmax');
    dicePixelClassification3dLayer('output');
];

将输入层和编码器模块与第四个解码器模块串联起来。将其他解码器模块作为单独的分支添加到层次图中。

layers = [layers; decoderModule4];
lgraph = layerGraph(layers);
lgraph = addLayers(lgraph,decoderModule3);
lgraph = addLayers(lgraph,decoderModule2);
lgraph = addLayers(lgraph,decoderModuleFinal);

使用串联层将每个编码器模块的第二个 reLu 层与解码器模块中同等大小的转置卷积层连接起来。将每个串联层的输出连接到解码器模块的第一个卷积层。

concat1 = concatenationLayer(4,2,'Name','concat1');
lgraph = addLayers(lgraph,concat1);
lgraph = connectLayers(lgraph,'en1_relu2','concat1/in1');
lgraph = connectLayers(lgraph,'de2_transconv','concat1/in2');
lgraph = connectLayers(lgraph,'concat1/out','de1_conv1');

concat2 = concatenationLayer(4,2,'Name','concat2');
lgraph = addLayers(lgraph,concat2);
lgraph = connectLayers(lgraph,'en2_relu2','concat2/in1');
lgraph = connectLayers(lgraph,'de3_transconv','concat2/in2');
lgraph = connectLayers(lgraph,'concat2/out','de2_conv1');

concat3 = concatenationLayer(4,2,'Name','concat3');
lgraph = addLayers(lgraph,concat3);
lgraph = connectLayers(lgraph,'en3_relu2','concat3/in1');
lgraph = connectLayers(lgraph,'de4_transconv','concat3/in2');
lgraph = connectLayers(lgraph,'concat3/out','de3_conv1'); 

您也可以使用 createUnet3d 辅助函数来创建三维 U-Net 层。此函数作为支持文件包含在本示例中。

lgraph = createUnet3d(inputSize);

绘制层次图。

analyzeNetwork(lgraph)

指定训练选项

使用 "adam" 优化求解器训练网络。使用 trainingOptions 函数指定超参数设置。初始学习率设置为 5e-4,并在训练期间逐渐降低。

options = trainingOptions('adam','MaxEpochs',100, ...
    'InitialLearnRate',5e-4, ...
    'LearnRateSchedule','piecewise', ...
    'LearnRateDropPeriod',5, ...
    'LearnRateDropFactor',0.95, ...
    'ValidationData',patchdsVal, ...
    'ValidationFrequency',400, ...
    'Plots','training-progress', ...
    'MiniBatchSize',miniBatchSize);

下载预训练网络和样本测试集

您也可以选择下载预训练的三维 U-Net 和 BraTS 数据集中的五个样本测试图像体及其对应的标签。预训练的模型和样本数据使您能够对测试数据执行分割,而无需下载完整的数据集或等待网络完成训练。

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

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

训练网络

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

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

注意:在 NVIDIA™ Titan X 上训练需要大约 60 个小时,根据您的 GPU 硬件情况,训练时间可能会更长。

doTraining = false;
if doTraining
    [net,info] = trainNetwork(augpatchds,lgraph,options);
else
    load(fullfile(imageDir,'trained3DUNet','brainTumor3DUNet.mat'));
end

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

对测试数据执行分割

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

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

通过使用辅助函数 centerCropMatReader,将图像和标签的中心部分裁剪为 128×128×128 体素大小。此函数作为支持文件包含在本示例中。voldsTest 变量存储真实值测试图像。pxdsTest 变量存储真实值标签。

windowSize = [128 128 128];
volReader = @(x) centerCropMatReader(x,windowSize);
labelReader = @(x) centerCropMatReader(x,windowSize);
voldsTest = imageDatastore(volLocTest, ...
    'FileExtensions','.mat','ReadFcn',volReader);
pxdsTest = pixelLabelDatastore(lblLocTest,classNames,pixelLabelID, ...
    'FileExtensions','.mat','ReadFcn',labelReader);

对于每个测试图像,将真实值图像体和标签添加到元胞数组。使用经过训练的网络和 semanticseg 函数来预测每个测试图像体的标签。

在执行分割后,通过将非大脑体素标记为 1(对应于背景)对预测的标签进行后处理。使用测试图像来确定哪些体素不属于大脑。您还可以通过使用 medfilt3 函数删除孤岛和填充孔来清理预测标签。medfilt3 不支持分类数据,因此在计算前将像素标签 ID 转换为 uint8。然后,将筛选后的标签转换回分类数据类型,为其指定原始像素标签 ID 和类名。

id=1;
while hasdata(voldsTest)
    disp(['Processing test volume ' num2str(id)])
    
    groundTruthLabels{id} = read(pxdsTest);
    
    vol{id} = read(voldsTest);
    tempSeg = semanticseg(vol{id},net);

    % Get the non-brain region mask from the test image.
    volMask = vol{id}(:,:,:,1)==0;
    % Set the non-brain region of the predicted label as background.
    tempSeg(volMask) = classNames(1);
    % Perform median filtering on the predicted label.
    tempSeg = medfilt3(uint8(tempSeg)-1);
    % Cast the filtered label to categorial.
    tempSeg = categorical(tempSeg,pixelLabelID,classNames);
    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 = 2;
vol3d = vol{volId}(:,:,:,1);

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

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

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

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

figure
h1 = labelvolshow(groundTruthLabels{volId},vol3d);
h1.LabelVisibility(1) = 0;
h1.VolumeThreshold = 0.68;

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

figure
h2 = labelvolshow(predictedLabels{volId},vol3d);
h2.LabelVisibility(1) = 0;
h2.VolumeThreshold = 0.68;

以下图像展示了逐层显示整个图像体切片的过程。

量化分割准确度

使用 dice 函数测量分割准确度。此函数计算预测和真实值分割之间的 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.99341
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.85204

下图显示一个 boxplot,它可视化包含五个样本测试四维体的集合中 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

摘要

此示例说明如何创建和训练一个三维 U-Net 网络来基于 BraTS 数据集执行三维脑肿瘤分割。训练网络的步骤包括:

  • 下载并预处理训练数据。

  • 定义一个 randomPatchExtractionDatastore 以向网络馈送训练数据。

  • 定义三维 U-Net 网络的各层。

  • 指定训练选项。

  • 使用 trainNetwork 函数训练网络。

在训练三维 U-Net 网络或加载预训练的三维 U-Net 网络后,该示例对测试数据集执行语义分割。该示例采用两种方法评估预测分割:其一,将预测分割与真实值分割进行可视化比较;其二,测量预测分割与真实值分割之间的 Dice 相似性系数。

参考

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

另请参阅

| | | | | |

相关示例

详细信息