CODEWITHSUNDEEP
kerasimage-segmentationlossunet-neural-network
user1147964
user1147964

Reputation: 143

Threshold the output of image segmentation of a keras model

I've implemented a U-Net in Keras for an image-segmentation problem I'm working on. I'm finding that the output in the early epochs provides very good results albeit with low probabilities. As the training continues and the IoU score increases, I get fewer pixels in my output, but they have higher probabilities.

To be more succinct, my output misses certain features but has higher probabilities in the features it does identify.

I'm curious to know if I can rectify this by forcing pixel values to be 1 after the output layer with the sigmoid activation function. Is there a simple way to do this? I'm using binary_crossentropy as the loss function so I was hoping I could do something like:

def binary_crossentropy_thresh(y_true,y_pred):
    y_pred[y_pred>=0.2] = 1

    return binary_crossentropy(y_true,y_pred)

But this won't work on Tensors. Any suggestions would be greatly appreciated.

Apologies if I haven't explained my problem sufficiently.

Upvotes: 1

Views: 778

Answers (1)

Daniel Möller
Daniel Möller

Reputation: 86610

No, you can't transform your function into a constant. It will not be differentiable (with respect to weights) and you will not be able to train.

You can do this in "custom metrics", never in the model or in the loss.

Make sure your IoU loss is being calculated "image-wise", not "batch-wise". The more general the loss, the less it will notice tiny things. Batchwise losses may ignore whole images if these images have less area than others.

A good suggestion of a loss function that worked great for me was a Lovasz loss.

You could also try sample weights (works best with single channel masks):

total_pixels = y_train.sum(axis=(1,2,3))
mean_pixels = total_pixels.mean(axis=0)
sample_weights = mean_pixels / total_pixels

Upvotes: 2

Related Questions