Defining threshold during image segmentation prediction

I am doing an image segmentation task where there are just two labels ‘0’ and 1(or black and white pixels).
I was wondering if there is a way to give a threshold when using unet_learner, so that a white pixel is chosen only if the value > threshold(say 0.6). I’m assuming by default it takes max of two possible outputs.


I had a similar question @sumeetd . Did you end up finding an answer to your question?

For what it’s worth, I found a nice post on Kaggle using a thresholding approach with binary classification, but I don’t see it build into the documentation.

Did either of you figure this out?

@aksg87 @jona Couldn’t find a solution yet.

You only need to use learn.predict and check third value in returned tuple which is the logits of probabilities. Then you use torch.sigmoid(x) > threshold to get your class prediction.