Semantic segmentation failing in small instance detection

Question

I performed semantic segmentation with U-net. My dataset consists of grayscale images of defects. After training the dataset for I got an metric accuracy of only 0.3 - 0.4 IOU. Eventhough it is merely low it performs well enough to identify instances that are huge means the prediction performs well enough in places where there is a standard intensity change(color change) and they are bigger instances. There are many other instances where there is no color change and it occoupies only few pixels in image(smaller instances) and the prediction rate is almost 0 on these instances.

I also tried Resdiual connection in the downsampling part of U-net likewise in ResNet. But still its the same and for smaller instances I used dilated convolution blocks in between the skip connections for encoder and decoder of U-net based on some papers. But still I cannot have a higher accuracy in my network and prediction rate for smaller instances are really poor. Although I use only 350 images for training with Data Augmentations. My image size is also 256,256.

Is there any other method I can try to increase the accuracy and prediction rate for smaller instances? Any suggestion would be helpful.

Answer 1

You may find it useful to categorize these smaller defects as a separate class, then introduce a class weights matrix to penalize incorrect classification of the smaller defects more heavily. If these small defects represent a very small portion of the total number of pixels in your training data, then the model may be stuck at a local minimum, where it just predicts them as zero because the loss is not heavily penalized for this. So you need to add extra penalty for the incorrect classification.

Assuming you create a new class for the small defects, the class weights matrix is implemented as a keras callback function.

If you do not want to create a multiclass problem, there is also this paper : https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5977596/#!po=32.4561 The multilevel unet architecture will not solve class imbalance if that is truly the underlying issue, but if may make the network more aware of smaller scale changes in the images you're working with.