search

NVIDIA / DIGITS

Deep Learning GPU Training System
https://developer.nvidia.com/digits
BSD 3-Clause "New" or "Revised" License
4.12k stars 1.38k forks source link

Segmentation with very unbalanced data #1172

Closed naranjuelo closed 7 years ago

naranjuelo commented 8 years ago

Hi! I read in this image segmentation tutorial that the solution to deal with very unbalanced datasets is to penalize the network when predicting some classes or to implement the Dice Python layer backward() method.

I have this problem when I try to do segmentation of small objects in big images (few pixels of the desired class comparing to the background). I was wondering if there is another more straightforward alternative in DIGITS, any way to normalize the number of pixels of each class in the error or another more appropriate loss function than SoftmaxWithLoss. Thank you!!

gheinrich commented 8 years ago

Hi @naranjuelo you could try cropping images so that the area to segment represents a bigger fraction of the overall image - though I imagine in most cases it is impractical to do so.

You could try using the InfoGainLoss layer to penalize more strongly wrong predictions of foreground. See this post. However the InfoGainLoss layer needs adjustments to work with FCNs - and to be honest when I tried it I didn't get the expected results.

I never tried implementing the backward route for the Dice metric but I believe this could be very efficient in those cases when you just have background and foreground.

One other thing you could try is to use Euclidean Loss - that should work well too if you only have background/foreground. To penalize bad foreground predictions more strongly I think you would just need to assign labels accordingly (e.g. background=1, foreground=10).

naranjuelo commented 8 years ago

Hi @gheinrich thank you for your suggestions. I discarded the cropping option, but maybe I could train a network with small crops to balance data and then apply it to the big image with no resizing (as you did in binary segmentation tutorial). I guess it could work.

About Euclidean Loss, if I work with only 2 classes, wouldn't it penalize any bad prediction in the same way? I mean, the euclidean distance between ground truth=1 and label=10 is the same as the distance when ground truth=10 and label=1. Maybe I misunderstood the function.

gheinrich commented 8 years ago

About Euclidean Loss, if I work with only 2 classes, wouldn't it penalize any bad prediction in the same way? I mean, the euclidean distance between ground truth=1 and label=10 is the same as the distance when ground truth=10 and label=1. Maybe I misunderstood the function.

Indeed the Euclidean distance is unsigned and you'll get the same result if you do foreground=1 and background=10 or background=1 and foreground=10 . So I think you're right, it won't penalize bad foreground predictions. It's worth trying though as EuclideanLoss might converge more easily and it is known for strongly penalizing outliers. Besides, the output will be a "score" and you can play with a threshold to determine where to draw the line between background and foreground.

naranjuelo commented 8 years ago

Okay, I'll give it a try then. Thanks!

gheinrich commented 7 years ago

@naranjuelo any update?

naranjuelo commented 7 years ago

Hi @gheinrich, I tried Euclidean Loss with no success. Now I'm working with multi-class segmentation (urban scenarios) and in this case the balancing problem doesn't happen even if data is not balanced (more classes, more images). The cropping solution is not very practical but should work (definitely), so when I have some time I'll prepare data and try it.

naranjuelo commented 7 years ago

It worked cropping patches where the target object represents a bigger fraction. Thanks!!