Merge pull request #128 from KatherLab/dev/logit-heatmap

Adjust heatmap function

Author: EzicStar

src/stamp/heatmaps/__init__.py

@@ -12,6 +12,7 @@
 from matplotlib.axes import Axes
 from matplotlib.figure import Figure
 from matplotlib.patches import Patch
+from packaging.version import Version
 from PIL import Image
 from torch import Tensor
 from torch.func import jacrev  # pyright: ignore[reportPrivateImportUsage]
@@ -22,8 +23,6 @@
 from stamp.preprocessing import supported_extensions
 from stamp.preprocessing.tiling import get_slide_mpp_
 from stamp.types import DeviceLikeType, Microns, SlideMPP, TilePixels
-from packaging.version import Version
-
 
 _logger = logging.getLogger("stamp")
 
@@ -33,25 +32,26 @@ def _gradcam_per_category(
     feats: Float[Tensor, "tile feat"],
     coords: Float[Tensor, "tile 2"],
 ) -> Float[Tensor, "tile category"]:
-    feat = -1  # feats dimension
+    feat_dim = -1
 
-    return (
+    cam = (
         (
             feats
             * jacrev(
-                lambda bags: torch.softmax(
-                    model.forward(
-                        bags=bags.unsqueeze(0),
-                        coords=coords.unsqueeze(0),
-                        mask=None,
-                    ),
-                    dim=1,
+                lambda bags: model.forward(
+                    bags=bags.unsqueeze(0),
+                    coords=coords.unsqueeze(0),
+                    mask=None,
                 ).squeeze(0)
             )(feats)
         )
-        .mean(feat)  # type: ignore
+        .mean(feat_dim)  # type: ignore
         .abs()
-    ).permute(-1, -2)
+    )
+
+    cam = torch.softmax(cam, dim=-1)
+
+    return cam.permute(-1, -2)
 
 
 def _vals_to_im(