🔧 Refactor heatmap activation and improve numerical stability in Soft…

…maxND

src/landmarker/heatmap/decoder.py

@@ -139,7 +139,6 @@ def _activate_norm_heatmap(
     if activation is not None:
         if activation == "softmax":
             heatmap = torch.exp(t * heatmap)
-            heatmap = heatmap / torch.sum(heatmap, dim=dim, keepdim=True)
         elif activation == "sigmoid":
             heatmap = torch.sigmoid(heatmap)
         elif activation == "ReLU":

src/landmarker/models/utils.py

@@ -8,7 +8,7 @@ def __init__(self, spatial_dims):
         self.dim = (-2, -1) if spatial_dims == 2 else (-3, -2, -2)
 
     def forward(self, x):
-        out = torch.exp(x)
+        out = torch.exp(x - torch.max(x, dim=self.dim, keepdim=True)[0])
         return out / torch.sum(out, dim=self.dim, keepdim=True)