딥러닝을 사용한 3차원 뇌종양 분할

사전 훈련된 3차원 U-Net 불러오기

사전 훈련된 3차원 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 데이터 세트에서 5개의 샘플 테스트 볼륨 및 그에 대응되는 레이블을 다운로드합니다. 이 헬퍼 함수는 예제에 지원 파일로 첨부되어 있습니다. 샘플 데이터를 사용하면 전체 데이터 세트를 다운로드하지 않고도 테스트 데이터에 대해 분할을 수행할 수 있습니다.

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 함수를 사용하여 각 블록의 레이블을 예측한 후 블록들을 완전한 분할된 테스트 볼륨으로 다시 결합합니다. 이 전략은 메모리 리소스가 제한된 GPU에서 효율적인 처리를 가능하게 합니다. 이 전략은 신경망에 있는 컨벌루션의 유효한 부분[5]을 사용하여 경계 아티팩트를 줄이기도 합니다.

볼륨 데이터를 blockedImage 객체로 저장하고 apply 함수를 사용해 블록을 처리하여 중첩 타일 전략을 구현합니다.

이전 섹션에서 다운로드한 샘플 볼륨에 대한 blockedImage 객체를 만듭니다.

bim = blockedImage(volTest);

apply 함수는 blockedImage 내의 각 블록에 대해 사용자 지정 함수를 실행합니다. semanticsegBlock을 각 블록에 대해 실행할 함수로 정의합니다.

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

블록 크기를 신경망 출력 크기로 지정합니다. 중첩하는 블록을 만들려면 0이 아닌 테두리 크기를 지정하십시오. 이 예제에서는 블록에 테두리를 더한 값이 신경망 입력 크기와 일치하도록 테두리 크기를 사용합니다.

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"가 적절합니다. 메모리 리소스가 제한된 GPU에서 메모리 부족 오류를 방지하기 위해 배치 크기는 1로 지정됩니다. 그러나 GPU에 충분한 메모리가 있는 경우 블록 크기를 늘려 처리 속도를 높일 수 있습니다.

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

실측 데이터의 중앙 단층과 깊이 방향의 예측 레이블을 보여주는 몽타주를 표시합니다.

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)")

다음 영상은 볼륨 중 하나에서 단층을 순차적으로 표시한 결과를 보여줍니다. 왼쪽은 레이블이 지정된 실측 데이터이고 오른쪽은 신경망 예측값입니다.

BraTS 데이터 세트 다운로드하기

훈련 데이터 세트 다운로드나 신경망 훈련을 원치 않는 경우에는 이 예제의 신경망 성능 평가하기 섹션으로 건너뛸 수 있습니다.

이 예제에서는 BraTS 데이터 세트 [2]를 사용합니다. BraTS 데이터 세트는 뇌종양 중에서도 가장 일반적인 주요 악성 뇌종양인 신경교종의 MRI 스캔을 포함합니다. 데이터 파일의 크기는 최대 7GB입니다.

BraTS 데이터를 다운로드하려면 Medical Segmentation Decathlon 웹사이트로 이동하여 "Download Data" 링크를 클릭하십시오. "Task01_BrainTumour.tar" 파일 [3]을 다운로드합니다. imageDir 변수로 지정된 디렉터리에 TAR 파일의 압축을 풉니다. 압축을 풀면 imageDir 안에 Task01_BrainTumour 디렉터리가 있고, 또 그 안에는 imagesTr, imagesTs, labelsTr이라는 3개의 하위 디렉터리가 있을 것입니다.

이 데이터 세트에는 3차원 영상이 쌓여있는 4차원 볼륨이 750개 있습니다. 각 4차원 볼륨은 크기가 240×240×155×4입니다. 여기서 처음 3개의 차원은 3차원 볼륨 영상의 높이, 너비, 깊이에 해당합니다. 4번째 차원은 각종 스캔 모달리티에 해당합니다. 데이터 세트는 복셀 레이블이 있는 훈련 볼륨 484개와 테스트 볼륨 266개로 나뉘어 있습니다. 테스트 볼륨에는 레이블이 없으므로 이 예제에서는 테스트 데이터를 사용하지 않습니다. 대신 이 예제에서는 484개의 훈련 볼륨을 각각 훈련, 검증, 테스트용으로 사용할 3개의 독립적인 세트로 분할합니다.

훈련을 위해 데이터 준비하기

3차원 U-Net 신경망을 더 효율적으로 훈련시키려면 헬퍼 함수 preprocessBraTSDataset을 사용하여 MRI 데이터를 전처리하십시오. 이 함수는 예제에 지원 파일로 첨부되어 있습니다. 헬퍼 함수는 다음과 같은 작업을 수행합니다.

데이터 전처리가 완료되는 데는 약 30분 정도 소요될 수 있습니다.

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

훈련 및 검증을 위해 무작위 패치 추출 데이터저장소 만들기

3차원 영상 데이터를 저장할 imageDatastore를 만듭니다. MAT 파일 형식은 비표준 영상 형식이므로 영상 데이터를 읽을 수 있도록 MAT 파일 리더를 사용해야 합니다. 헬퍼 MAT 파일 리더 matRead를 사용하면 됩니다. 이 함수는 예제에 지원 파일로 첨부되어 있습니다.

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

레이블을 저장할 pixelLabelDatastore를 만듭니다.

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

실측 영상과 그에 해당하는 픽셀 레이블 데이터로부터 무작위 패치를 추출하는 randomPatchExtractionDatastore를 만듭니다. 패치 크기를 132×132×132 복셀로 지정합니다. 훈련 중에 각 볼륨 및 레이블 쌍으로부터 무작위로 배치된 16개의 패치를 추출하도록 "PatchesPerImage"를 지정합니다. 미니 배치 크기를 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;

3차원 U-Net 구성하기

이 예제에서는 3차원 U-Net 신경망 [1]을 사용합니다. U-Net에서는 초기 일련의 컨벌루션 계층의 사이에 일정 간격으로 최댓값 풀링 계층이 배치되어 입력 영상의 해상도를 연이어 낮춥니다. 그다음에 오는 일련의 컨벌루션 계층의 사이에는 일정 간격으로 업샘플링 연산자가 배치되어 입력 영상의 해상도를 연이어 높입니다. 각 ReLU 계층 전에는 배치 정규화 계층이 옵니다. U-Net이라는 이름은 이 신경망을 문자 U와 같은 대칭 형태로 그릴 수 있다는 사실에서 유래했습니다.

unetLayers 함수를 사용하여 디폴트 3차원 U-Net 신경망을 만듭니다. 2개 클래스 분할을 지정합니다. 테스트 볼륨 예측을 위해 중첩 타일 전략을 사용할 때 경계 아티팩트를 방지하기 위해 유효한 컨벌루션 채우기도 지정합니다.

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

transform 함수를 헬퍼 함수 augmentAndCrop3dPatch로 지정된 사용자 지정 전처리 연산과 함께 사용하여 훈련 데이터와 검증 데이터를 증대합니다. 이 함수는 예제에 지원 파일로 첨부되어 있습니다. augmentAndCrop3dPatch 함수는 다음 연산을 수행합니다.

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

이 예제에서는 작은 종양 영역을 더 잘 분할하고 큰 배경 영역의 영향을 줄이기 위해 dicePixelClassificationLayer를 사용합니다. 픽셀 분류 계층을 다이스 픽셀 분류 계층으로 교체합니다.

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

데이터는 이 예제의 훈련 데이터와 검증 데이터 전처리하기 섹션에서 이미 정규화되었습니다. image3dInputLayer (Deep Learning Toolbox)에서의 데이터 정규화는 불필요하므로 입력 계층을 데이터 정규화가 없는 입력 계층으로 교체합니다.

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

또는 심층 신경망 디자이너 앱을 사용하여 3차원 U-Net 신경망을 수정할 수도 있습니다.

deepNetworkDesigner(lgraph)

3차원 U-Net 훈련시키기

훈련 옵션 지정하기

adam 최적화 솔버를 사용하여 신경망을 훈련시킵니다. trainingOptions (Deep Learning Toolbox) 함수를 사용하여 하이퍼파라미터 설정을 지정합니다. 초기 학습률은 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);

신경망 훈련시키기

기본적으로 이 예제에서는 다운로드한 사전 훈련된 3차원 U-Net 신경망을 사용합니다. 사전 훈련된 신경망을 사용하면 훈련이 완료될 때까지 기다리지 않고 의미론적 분할을 수행하고 분할 결과를 평가할 수 있습니다.

신경망을 훈련시키려면 다음 코드에서 doTraining 변수를 true로 설정하십시오. trainNetwork (Deep Learning Toolbox) 함수를 사용하여 신경망을 훈련시킵니다.

사용 가능한 GPU가 있으면 GPU에서 훈련시킵니다. GPU를 사용하려면 Parallel Computing Toolbox™와 CUDA® 지원 NVIDIA® GPU가 필요합니다. 자세한 내용은 GPU 연산 요구 사항 (Parallel Computing Toolbox) 항목을 참조하십시오. 훈련은 4개의 NVIDIA™ Titan Xp 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

3차원 U-Net 평가하기

테스트를 위한 실측 볼륨과 레이블을 포함하는 테스트 데이터 소스를 선택합니다. 다음 코드에서 useFullTestSet 변수를 false로 둘 경우 이 예제에서는 테스트를 위해 5개의 샘플 볼륨을 사용합니다. 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 함수를 사용하여 각 블록을 처리합니다. 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 함수를 사용하여 분할에 대한 데이터 세트 메트릭 및 블록 메트릭을 평가합니다.

[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  

각 영상에 대해 계산된 자카드 점수를 표시합니다.

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/

The BraTS dataset is provided by Medical Segmentation Decathlon under the CC-BY-SA 4.0 license. All warranties and representations are disclaimed; see the license for details. MathWorks® has modified the data set linked in the Download BraTS Sample Data section of this example. The modified sample data set has been cropped to a region containing primarily the brain and tumor and each channel has been normalized independently by subtracting the mean and dividing by the standard deviation of the cropped brain region.

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