Unet網路

Unet是一種編碼-解碼結構相結合的神經網路結構，是一種語義分割網路。在醫學影象分割的相關應用中被廣泛使用。使用matlab可以快速實現網路結構的定義和訓練。

資料集準備

準備待訓練影象和相對應的標註影象，將影象和標註影象分別存放到不同的目錄中，通過相同的檔名進行一一對應。

%% 資料集載入
dataSetDir = fullfile('./data');
imageDir = fullfile(dataSetDir,'trainingImages');
labelDir = fullfile(dataSetDir,'trainingLabels');

定義畫素分類的類別名稱，以及各類別在標註影象中的亮度值

classNames = ["triangle","background"];
labelIDs  = [255 0];

生成訓練資料集物件

imds = imageDatastore(imageDir);
pxds = pixelLabelDatastore(labelDir,classNames,labelIDs);
% ds = pixelLabelImageDatastore(imds,pxds);
ds = combine(imds,pxds);

網路定義

imageSize = [32 32];
numClasses = 2;
lgraph = unetLayers(imageSize, numClasses)
 

訓練網路

options = trainingOptions('sgdm', ...
  'InitialLearnRate',1e-3, ...
  'MaxEpochs',20, ...
  'VerboseFrequency',10);

net = trainNetwork(ds,lgraph,options)

匯出ONNX格式的模型，可使用opencv或tensorrt等工具進行應用部署

exportONNXNetwork(net,'myunet.onnx');

測試

pic = imread('.\data\testImages\image_002.jpg');
out2 = predict(net,pic);

subplot(1,2,1)
imshow(pic)
subplot(1,2,2)
imshow(out2(:,:,1))
 

完成程式碼和測試資料

https://download.csdn.net/download/Ango_/16138054