I haven’t studied the architecture of MobileNetv3, but to generate heatmaps(or in general activation maps as such), a fully convolutional network should suffice where there is a feature extraction stem followed by a 1x1 conv to bring down the number of channels to 1 keeping the spatial dims as it is.
Coming to the problem of class imbalance, you can counter it using Balanced/Weighted Cross Entropy or Dice loss. These loss functions are generally used to tackle class imbalance in detection/segmentation tasks.