Unet network
Unet is a neural network structure that combines an encoding-decoding structure, and is a semantic segmentation network. It is widely used in related applications of medical image segmentation. Using matlab can quickly realize the definition and training of network structure.
Dataset Preparation
Prepare the image to be trained and the corresponding labeled image, store the image and the labeled image in different directories, and use the same file name for one-to-one correspondence.
%% 数据集加载
dataSetDir = fullfile('./data');
imageDir = fullfile(dataSetDir,'trainingImages');
labelDir = fullfile(dataSetDir,'trainingLabels');Define the category name of the pixel classification, and the brightness value of each category in the labeled image
classNames = ["triangle","background"];
labelIDs = [255 0];Generate training dataset object
imds = imageDatastore(imageDir);
pxds = pixelLabelDatastore(labelDir,classNames,labelIDs);
% ds = pixelLabelImageDatastore(imds,pxds);
ds = combine(imds,pxds);network definition
imageSize = [32 32];
numClasses = 2;
lgraph = unetLayers(imageSize, numClasses)train the network
options = trainingOptions('sgdm', ...
'InitialLearnRate',1e-3, ...
'MaxEpochs',20, ...
'VerboseFrequency',10);
net = trainNetwork(ds,lgraph,options)Export the model in ONNX format, and use tools such as opencv or tensorrt for application deployment
exportONNXNetwork(net,'myunet.onnx');test
pic = imread('.\data\testImages\image_002.jpg');
out2 = predict(net,pic);
subplot(1,2,1)
imshow(pic)
subplot(1,2,2)
imshow(out2(:,:,1))
Complete code and test data
https://download.csdn.net/download/Ango_/16138054