# Semantic segmentation CNN outputs all zeroes

I'm using MATLAB 2019, Linux, and UNet (a CNN specifically designed for semantic segmentation). I'm training the network to classify all pixels in an image as either cell or background to get segmentations of cells in microscopic images. My problem is the network is classifying every single pixel as background, and seems to just be outputting all zeroes. The validation accuracy improves a little at the very start of the training but than plateaus at around 60% for the majority of the training time. The network doesn't seem to be training very well and I have no idea why.

Can anyone give me some hints about what I should look into more closely? I just don't even know where to start with debugging this.

Here's my code:

    % Set datapath
datapath = '/scratch/qbi/uqhrile1/ethans_lab_data';

% Get training and testing datasets
images_dataset = imageDatastore(strcat(datapath,'/bounding_box_cropped_resized_rgb'));
labels = pixelLabelDatastore(gTruth);
[imdsTrain, imdsVal, imdsTest, pxdsTrain, pxdsVal, pxdsTest] = partitionCamVidData(images_dataset,labels);

% Weight segmentation class importance by the number of pixels in each class
pixel_count = countEachLabel(labels); % count number of each type of pixel
frequency = pixel_count.PixelCount ./ pixel_count.ImagePixelCount; % calculate pixel type frequencies
class_weights = mean(frequency) ./ frequency; % create class weights that balance the loss function so that more common pixel types won't be preferred

% Specify the input image size.
imageSize = [512 512 3];

% Specify the number of classes.
numClasses = 2;

% Create network
lgraph = unetLayers(imageSize,numClasses);

% Replace the network's classification layer with a pixel classification
% layer that uses class weights to balance the loss function
pxLayer = pixelClassificationLayer('Name','labels','Classes',pixel_count.Name,'ClassWeights',class_weights);
lgraph = replaceLayer(lgraph,"Segmentation-Layer",pxLayer);

%% TRAIN THE NEURAL NETWORK

% Define validation dataset-with-labels
validation_dataset_with_labels = pixelLabelImageDatastore(imdsVal,pxdsVal);

% Training hyper-parameters: edit these settings to fine-tune the network
options = trainingOptions('adam', 'LearnRateSchedule','piecewise', 'LearnRateDropPeriod',10, 'LearnRateDropFactor',0.3, 'InitialLearnRate',1e-3, 'L2Regularization',0.005, 'ValidationData',validation_dataset_with_labels, 'ValidationFrequency',10, 'MaxEpochs',3, 'MiniBatchSize',1, 'Shuffle','every-epoch');

% Set up data augmentation to enhance training dataset
aug_imgs = {};
numberOfImages = length(imdsTrain.Files);
for k = 1 : numberOfImages
% Apply cutout augmentation
cutout_img = random_cutout(img);imwrite(cutout_img,strcat('/scratch/qbi/uqhrile1/ethans_lab_data/augmented_dataset/img_',int2str(k),'.tiff'));
end
aug_imdsTrain = imageDatastore('/scratch/qbi/uqhrile1/ethans_lab_data/augmented_dataset');
augmenter = imageDataAugmenter('RandXReflection',true, 'RandXTranslation',[-10 10],'RandYTranslation',[-10 10]);
% Combine augmented data with training data
augmented_training_dataset = pixelLabelImageDatastore(aug_imdsTrain, pxdsTrain, 'DataAugmentation',augmenter);

% Train the network
[cell_segmentation_nn, info] = trainNetwork(augmented_training_dataset,lgraph,options);

save cell_segmentation_nn


Similar to other answers, I don't know Matlab that well but you could try the following steps to debug your problem.

# Make sure you can overfit to a single instance

from your dataset, pull out a single image with a good amount of true positives in it. Duplicate that images B times (where B = Batch Size) and then try to train your network with only that small dataset. If you can't overfit to a single instance, then something is really wrong and you should validate all functional aspects of your network. If you can overfit, then it's probably more of an algorithmic or data imbalance issue.

# Validate Functional Aspects of the Network

### Make sure your input data and labels are correct.

Validate that your images are being correclty inputted into the network. You can do this by printing out any images right before they go into the training function. For the labels, make sure to manually inspect a few labels to be sure that the labels correctly match up the images.

### Make sure that your loss function is correct.

Add a unit test or two to your loss function to validate that it is doing what it should be doing. Create a simiple example that you can easily validate.

### Validate the rest of the non-algorithmic functionality

Anything that isn't a design choice should be validated. Make sure that the weights are the correct sizes, each layer has the correct number of weights, the intermediate features have the correct shape, etc.

# Data Imbalance

If you were able to overfit to a single image, and validated the functional aspects of your network then you may be facing a data imbalance issue. Look at your dataset and see what % of your instances are true positives vs. true negatives. If you have an exteme imbalance (like 10% 90%) or something like that, then build a dataset that is more balanced and see if you can fit your data. If you fit the data with that more balanced dataset, then there's plenty of ways to fix your data imbalance issue. Google around for data imbalance and you should get a few good ideas. Some include focal loss, upsampling, etc...