Hello,
As per my understanding, you have created a U-Net model using the below mentioned document as reference. Now you wish to create a multi-class U-Net model and split the output mask for each class.
As the mentioned documentation provides a simple convolution neural network for regression which outputs a single continuous value, I assume you have added up-sampling layers and skip-connections to make the network into U-Net architecture. To modify the network into multi-class U-Net model, consider changing the the last few layers after the last “relu” layer as show in example code below:
conv = convolution2dLayer(1, numClasses); % numClasses is the number of output classes
softmax = softmaxLayer();
outputLayer = pixelClassificationLayer();
After getting the output mask (“labelled_mask”) from the U-Net model, each class output can be separated using the example code below.
% Extract each class into a separate image
num_classes = 4;
for i = 1:num_classes
% Create binary mask for current class
class_mask = (labeled_mask == i);
% Apply binary mask to input image
class_img = img .* uint8(class_mask);
% Show output image
figure,
imshow(class_img);
end
Please refer to below MathWorks documentation for more details about the functions used:
- “softmaxLayer” function: https://www.mathworks.com/help/deeplearning/ref/nnet.cnn.layer.softmaxlayer.html
- “pixelClassificationLayer” function: https://www.mathworks.com/help/vision/ref/nnet.cnn.layer.pixelclassificationlayer.html
I hope this helps!
