使用深度学习进行三维脑肿瘤分割
此示例说明如何基于三维医学图像执行脑肿瘤的语义分割。
语义分割中,图像中的每个像素或三维体的每个体素都被标注为某一类。此示例说明如何在磁共振成像 (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)")
以下图像顺序显示了该三维体的各切片。标注的真实值在左侧,网络预测在右侧。
训练三维 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
的目录,该目录有三个子目录:imagesTr
、imagesTs
和 labelsTr
。
该数据集包含 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
函数执行以下操作:
随机旋转和翻转训练数据,使训练更加稳健。该函数不旋转或翻转验证数据。
将响应补片裁剪为网络的输出大小 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.
另请参阅
randomPatchExtractionDatastore
(Image Processing Toolbox) | trainNetwork
| trainingOptions
| transform
| pixelLabelDatastore
(Computer Vision Toolbox) | imageDatastore
| semanticseg
(Computer Vision Toolbox) | dicePixelClassificationLayer
(Computer Vision Toolbox)