Hi Wan!
GradCam is typically used for interpreting the predictions of convolutional neural networks (CNNs) on image classification tasks. However, your use case involves a U-Net architecture, which is commonly used for semantic segmentation tasks. This might be the reason why you are not getting the expected localization for your classes.
The U-Net architecture is different from the typical CNNs used for image classification. It consists of a contracting path (encoder), a bottleneck, and an expansive path (decoder).
In U-Net, the last layers of the decoder are responsible for upscaling the feature maps and making dense predictions for every pixel. You should choose a feature layer from the decoder that is close to the output layer but before the final classification layer. It will ensure that the feature maps are detailed enough for localization but also contain enough high-level information to represent the classes.
Hope it helps!
