0
$\begingroup$

I am trying to perform segmentation of coils on PC motherboards:

Example input:

enter image description here

Example output:

enter image description here

I have around 150 training samples; some have more coils or bigger coils and some have less, but on every image there is at least one coil.

I tried to realize the segmentation with a VGG-net that got a fully convolutional network on top of it. Without the fully connected layers of the VGG. Before that, I tried something similar with a way smaller architecture with the same outcome.

What I realized in both cases is that the net pretty fast starts to predict that the whole images is background. I thought now that this is the case because the coils are only a very small part of the whole image and for the net it is easier to optimize in the direction of an all black image.

I use as criterion BCEWithLogitsLoss and as optimizer RMSprop. Bellow are my statistics with the vggnet+fcn. In the IoUs array is the first value for the coils and the second for the background.(Source code with the vgg+fcn part from for my second try). I used vgg11 and just the normal fcn.

enter image description here

Can't I realize a segmentation with that dataset? Is the dataset too small? What else could be a reasons for this results?

Improve this question
$\endgroup$
2
  • $\begingroup$ Could you maybe add some better description of the architecture you are using? What do you mean by "VGG-net that got a fully convolutional network on top of it"? $\endgroup$
    – Jan Kukacka
    Jun 4, 2018 at 15:40
  • $\begingroup$ I used the VGGNet11 structure with 8 weighted layers because I cut of the 3 fully connected layers I don't need. The fully convolutional network who follows the 8 weighted layers contains of 5 deconv layers with batchnorm layers between them and one conv layer in the the end. If you check this github.com/pochih/FCN-pytorch/blob/master/python/fcn.py you can see the structure of the FCN part I mention $\endgroup$
    – Hille
    Jun 4, 2018 at 15:57

1 Answer 1

Reset to default
1
$\begingroup$

This problem can certainly be easily solved using convolutional neural networks. Actually, this is what they are pretty good at. Few possible reasons why your network does not work come to my mind:

  • Try using weighted loss binomial cross entropy function. Segmentation learning may easily fail when one class occupies larger portions of the image, compared to the other one. Alternatively, try using IoU or Dice coefficient as loss function.

  • 150 images is pretty small dataset. Make sure you use at least a pre-trained network, and even then consider getting more training data. Also, make sure you use data augmentation: It is a cheap way to get a lot more data.

Improve this answer
$\endgroup$
3
  • $\begingroup$ I used the pre-trained feature extracting part of the VGGNet. I thought about training my own VGGNet who does a binary classification of coils and background. Yeah I will try it with some augmentation too but after the first results with the 150 images I thought there must be something wrong with my approach because its convergence that fast to an all black image. $\endgroup$
    – Hille
    Jun 4, 2018 at 15:35
  • $\begingroup$ Try the weighted loss or dice. It really helps. $\endgroup$
    – Jan Kukacka
    Jun 4, 2018 at 15:39
  • $\begingroup$ I tried now to implement the weighted loss but I could not rly find a Pytorch function who did that the way I want it to. But I found an implementation of a dice function. Now I need the output of my net to be, BatchxnclassesxHxW representing log probabilities for each class. At the moment the output of my net is BatchxnclassesxHxW without log probabilities. Is there a function out there that can calculate the log probabilities? I tried the softmax function but that does not work with the input BatchxnclassesxHxW . $\endgroup$
    – Hille
    Jun 6, 2018 at 7:21

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.