Today I've made another test with a dataset consisting of 4 images edited through the median filter.
The net used just pure BCE as the loss function. Only for a single one of these images the net predicted the proper mask. My assumption is that the network somehow stops training after it sees the initial example. I will try to remove the shuffle() in train_epoch() and see whether this would cause the net to predict perfectly only the first example from the list of training files.
If this turned out to be random then that would mean that we're dealing with some crazy bad generalization problem.
I've run the training after removing the shuffles. The trained net made perfect prediction only for the second image in the training set.
first prediction
second prediction
This makes me think that there is definitely a generalization problem bothering us. This looks like it has to do with the discrepancy in count of mask pixels vs background pixels. On the initial image the net overshot the ground truth mask by a considerable margin. On that image the mask region was really small in comparison to overall image size. So this looks like a sort of broader focusing problem.
Let's maybe try to run the training on the same set of images with focal loss function this time. If that doesn't help we can try to pivot to something like the Attention U-Net.
Tried focal loss with multiple alpha-gamma combinations. The length of training was 100 epochs. The ones that look promising so far:
alpha = 0.6, gamma = 0.7
alpha = 0.9, gamma = 2
alpha = 0.9, gamma = 3
alpha = 0.99, gamma = 3.5
alpha = 0.999, gamma = 4
Today I'm experimenting with adding different loss functions to the mix. Turns out that adding in BCE actually degrades the training quality. I've tried two different proportions: 70% focal + 30% BCE and 80% focal + 20% BCE.
focal + bce
I also managed to try adding continuous DICE to the mix. It didin't help. It actually peformed slightly worse than focal alone on the example with the smallest tumor.
focal + continuous DICE
A thought came to my mind that displaying the network's output as a probability map with different colours representing different probabilities might be also helpful for debugging. I tried that and the prediction don't really look that different. Only the areas that were depicted in blue before have noticably higher probabilities than the rest of the image.
focal + continuous DICE with probabilities
Next I would like try whether adding Tversky loss to the mix could help.