I'm using U-Net for image segmentation.

The model was trained with images that could contain up to 4 different classes. The train classes are never overlapping.

The output of the UNet is a heatmap (with float values between 0 and 1) for each of these 4 classes.

Now, I have 2 problems:

  • for a certain class, how do I segment (draw contours) in the original image only for the points where the heatmap has significant values? (In the image below an example: the values in the centre are significant, while the values on the left aren't. If I draw the segmentation of the entire image without any additional operation, both are considered.)

enter image description here

  • downstream of the first point, how do I avoid that in the original image the contours of two superimposed classes are drawn? (maybe by drawing only the one that has higher values in the corresponding heatmap)
  • $\begingroup$ Regarding problem 1 (if I understand it correctly), have you tried something like non-maximum suppression, or something similar? $\endgroup$ – nbro Jul 5 at 23:34

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.