changes to gradcam

app.py

@@ -51,7 +51,6 @@ def root():
     </html>
     """
 
-
 @app.get("/health")
 def health():
     return {"status": "ok", "device": DEVICE, "model": "DenseNet121-CBAM"}

inference.py

@@ -97,17 +97,17 @@ class GradCAMPlusPlus:
         for name, module in self.model.named_modules():
             if name == target_layer_name:
                 self.hooks.append(module.register_forward_hook(forward_hook))
-                self.hooks.append(module.register_backward_hook(backward_hook))
+                self.hooks.append(module.register_full_backward_hook(backward_hook))
                 return
         raise ValueError(f"Layer '{target_layer_name}' not found in model")
 
     def generate_cam(self, input_tensor, class_idx=None):
         self.model.eval()
-        input_tensor = input_tensor.clone().detach().requires_grad_(True)
 
         for param in self.model.parameters():
             param.requires_grad = True
 
+        input_tensor = input_tensor.clone().detach().requires_grad_(True)
         output = self.model(input_tensor)
 
         if class_idx is None:
@@ -117,18 +117,34 @@ class GradCAMPlusPlus:
         self.model.zero_grad()
         target_score.backward(retain_graph=True)
 
-        grads = self.gradients        # [1, C, H, W]
-        acts  = self.activations      # [1, C, H, W]
+        grads = self.gradients    # [1, C, H, W]
+        acts  = self.activations  # [1, C, H, W]
+
+        B, C, H, W = grads.shape
 
-        alpha_num   = grads.pow(2)
-        alpha_denom = 2.0 * grads.pow(2) + \
-                      acts.mul(grads.pow(3)).view(1, grads.size(1), -1).sum(dim=-1, keepdim=True).unsqueeze(-1)
-        alpha_denom = torch.where(alpha_denom != 0.0, alpha_denom, torch.ones_like(alpha_denom))
-        alpha       = alpha_num.div(alpha_denom + 1e-7)
-        weights     = alpha.mul(torch.relu(grads).view(1, grads.size(1), -1).sum(dim=-1))
+        # ── GradCAM++ alpha computation (all ops stay in [B, C, H, W]) ──
+        grads_sq  = grads.pow(2)                          # [1, C, H, W]
+        grads_cub = grads.pow(3)                          # [1, C, H, W]
 
-        cam = (weights.unsqueeze(-1).unsqueeze(-1) * acts).sum(dim=1, keepdim=True)
-        cam = torch.relu(cam).squeeze().cpu().detach().numpy()
+        # sum over spatial dims H,W → [1, C, 1, 1]  then broadcast back
+        spatial_sum = (acts * grads_cub).sum(dim=[2, 3], keepdim=True)  # [1, C, 1, 1]
+
+        alpha_denom = 2.0 * grads_sq + spatial_sum        # [1, C, H, W]
+        alpha_denom = torch.where(
+            alpha_denom != 0.0,
+            alpha_denom,
+            torch.ones_like(alpha_denom)
+        )
+        alpha = grads_sq / (alpha_denom + 1e-7)           # [1, C, H, W]
+
+        # weights: alpha * ReLU(grads), summed over H,W → [1, C, 1, 1]
+        weights = (alpha * torch.relu(grads)).sum(dim=[2, 3], keepdim=True)  # [1, C, 1, 1]
+
+        # CAM: weighted sum of activations → [1, 1, H, W]
+        cam = (weights * acts).sum(dim=1, keepdim=True)   # [1, 1, H, W]
+        cam = torch.relu(cam)
+
+        cam = cam.squeeze().cpu().detach().numpy()         # [H, W]
         cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
         return cam
 
@@ -137,6 +153,7 @@ class GradCAMPlusPlus:
             hook.remove()
 
 
+
 # ─── Preprocessing ─────────────────────────────────────────────────────────────
 
 def get_transform(image_size=512):