Is it possible to calculate class prediction with heatmap?

[Makhaon]

For first, thank you for great realisation. I have a question.
I want to calulate not only heatmap but to get classes probabilities also. I want to do it in the HeatmapGenerator.py module for only one image (as it works now for the heatmap). Maybe somebody know Is it a way to do it? Thank you.