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.)
- 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)