When training a neural network does every training image need to have every pixel class?

I am trying to train a neural network that is able to take micrographs of polished cross-sections of polymer containing fibres and particles. I want the trained neural network to identiffy the pixels that correspond to the five different features: polymer, fibres, particle type A, particle type B, and background (black eges of images).

I have 126 training images which I have manually segmented. Below shows the original image (left), the normalised ground truth (centre), and the real ground truth (right).

I have a training image datastore (dsTrain) and I have set up the network as follows:

lgraph = unetLayers([296 296 1], 5,'encoderDepth',3);
options = trainingOptions('adam',...
'InitialLearnRate', .0001,...
'MiniBatchSize', 4,...
'MaxEpochs', 50,...
'ValidationData', dsVal,...
'Verbose', false,...
'ExecutionEnvironment', 'gpu',...
Plots = "training-progress");
net4 = trainNetwork(dsTrain,lgraph_1,options);

My questions are:

1. Does every training image have to contain pixels corresponding to each of the five features?
2. My ground truth is made up of images where the pixel values are as follows:

- Polymer = 0

- Fibres = 1

- Particle A = 2

- Particle B = 3

- Background = 4

Is this the best way to set this up?

Currently when I train the network the RMSE remains very high (> 500) after several hundred iterations.

Thank you in advance!

Robin

Answer 1

[Joss Knight]

The answer to question 1 is no. The answer to question 3 is that probably 126 training images is not sufficient, you might need about 10x that amount plus some data augmentation.

I don't know the answer to question 2.

[Robin Hartley]

Hi Joss,

Thank you for the response. It turns out the primary problem was that I had set up my network for image to image regression, when really I needed image to image classification.

I have also increased the number of training images and augmented them too.

Best wishes,

Robin