Matlab图像分割(U-Net)

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