Attempt to change all Numpy calls to Torch calls

pytorch_grad_cam/base_cam.py

@@ -36,7 +36,7 @@ def get_cam_weights(self,
                         target_layers: List[torch.nn.Module],
                         targets: List[torch.nn.Module],
                         activations: torch.Tensor,
-                        grads: torch.Tensor) -> np.ndarray:
+                        grads: torch.Tensor) -> torch.Tensor:
         raise Exception("Not Implemented")
 
     def get_cam_image(self,
@@ -45,7 +45,7 @@ def get_cam_image(self,
                       targets: List[torch.nn.Module],
                       activations: torch.Tensor,
                       grads: torch.Tensor,
-                      eigen_smooth: bool = False) -> np.ndarray:
+                      eigen_smooth: bool = False) -> torch.Tensor:
 
         weights = self.get_cam_weights(input_tensor,
                                        target_layer,
@@ -62,7 +62,7 @@ def get_cam_image(self,
     def forward(self,
                 input_tensor: torch.Tensor,
                 targets: List[torch.nn.Module],
-                eigen_smooth: bool = False) -> np.ndarray:
+                eigen_smooth: bool = False) -> torch.Tensor:
 
         if self.cuda:
             input_tensor = input_tensor.cuda()
@@ -73,7 +73,7 @@ def forward(self,
 
         outputs = self.activations_and_grads(input_tensor)
         if targets is None:
-            target_categories = np.argmax(outputs.cpu().data.numpy(), axis=-1)
+            target_categories = torch.argmax(outputs.data, axis=-1)
             targets = [ClassifierOutputTarget(
                 category) for category in target_categories]
 
@@ -106,10 +106,10 @@ def compute_cam_per_layer(
             self,
             input_tensor: torch.Tensor,
             targets: List[torch.nn.Module],
-            eigen_smooth: bool) -> np.ndarray:
-        activations_list = [a.cpu().data.numpy()
+            eigen_smooth: bool) -> torch.Tensor:
+        activations_list = [a.data
                             for a in self.activations_and_grads.activations]
-        grads_list = [g.cpu().data.numpy()
+        grads_list = [g.data
                       for g in self.activations_and_grads.gradients]
         target_size = self.get_target_width_height(input_tensor)
 
@@ -130,24 +130,24 @@ def compute_cam_per_layer(
                                      layer_activations,
                                      layer_grads,
                                      eigen_smooth)
-            cam = np.maximum(cam, 0)
+            cam = torch.maximum(cam, torch.tensor(0))
             scaled = scale_cam_image(cam, target_size)
             cam_per_target_layer.append(scaled[:, None, :])
 
         return cam_per_target_layer
 
     def aggregate_multi_layers(
             self,
-            cam_per_target_layer: np.ndarray) -> np.ndarray:
-        cam_per_target_layer = np.concatenate(cam_per_target_layer, axis=1)
-        cam_per_target_layer = np.maximum(cam_per_target_layer, 0)
-        result = np.mean(cam_per_target_layer, axis=1)
+            cam_per_target_layer: torch.Tensor) -> torch.Tensor:
+        cam_per_target_layer = torch.cat(cam_per_target_layer, axis=1)
+        cam_per_target_layer = torch.maximum(cam_per_target_layer, torch.tensor(0))
+        result = torch.mean(cam_per_target_layer, axis=1)
         return scale_cam_image(result)
 
     def forward_augmentation_smoothing(self,
                                        input_tensor: torch.Tensor,
                                        targets: List[torch.nn.Module],
-                                       eigen_smooth: bool = False) -> np.ndarray:
+                                       eigen_smooth: bool = False) -> torch.Tensor:
         transforms = tta.Compose(
             [
                 tta.HorizontalFlip(),
@@ -167,18 +167,18 @@ def forward_augmentation_smoothing(self,
             cam = transform.deaugment_mask(cam)
 
             # Back to numpy float32, HxW
-            cam = cam.numpy()
+            # cam = cam.numpy()
             cam = cam[:, 0, :, :]
-            cams.append(cam)
+            cams.append(cam) # TODO: Handle this for torch tensors
 
-        cam = np.mean(np.float32(cams), axis=0)
+        cam = torch.mean(cams.to(torch.float32), axis=0)
         return cam
 
     def __call__(self,
                  input_tensor: torch.Tensor,
                  targets: List[torch.nn.Module] = None,
                  aug_smooth: bool = False,
-                 eigen_smooth: bool = False) -> np.ndarray:
+                 eigen_smooth: bool = False) -> torch.Tensor:
 
         # Smooth the CAM result with test time augmentation
         if aug_smooth is True:

pytorch_grad_cam/grad_cam.py

@@ -1,4 +1,4 @@
-import numpy as np
+import torch
 from pytorch_grad_cam.base_cam import BaseCAM
 
 
@@ -19,4 +19,4 @@ def get_cam_weights(self,
                         target_category,
                         activations,
                         grads):
-        return np.mean(grads, axis=(2, 3))
+        return torch.mean(grads, axis=(2, 3))

pytorch_grad_cam/utils/image.py

@@ -4,7 +4,7 @@
 import cv2
 import numpy as np
 import torch
-from torchvision.transforms import Compose, Normalize, ToTensor
+from torchvision.transforms import Compose, Normalize, ToTensor, Resize
 from typing import List, Dict
 import math
 
@@ -158,16 +158,30 @@ def show_factorization_on_image(img: np.ndarray,
 
 
 def scale_cam_image(cam, target_size=None):
-    result = []
-    for img in cam:
-        img = img - np.min(img)
-        img = img / (1e-7 + np.max(img))
-        if target_size is not None:
-            img = cv2.resize(img, target_size)
-        result.append(img)
-    result = np.float32(result)
+    # Disabled the target_size scaling for now
+    # It appears to swap the axes dimensions and needs further work for the
+    # proof of concept
 
-    return result
+    # if target_size is not None:
+    #     result = torch.zeros([cam.shape[0], target_size[0], target_size[1]])
+    # else:
+    #     result = torch.zeros(cam.shape)
+
+    result = torch.zeros(cam.shape)
+
+    for i in range(cam.shape[0]):
+        img = cam[i]
+        img = img - torch.min(img)
+        img = img / (1e-7 + torch.max(img))
+
+        # if target_size is not None:
+            # transform = Resize(target_size)
+            # img = Resize(size = target_size)(img)
+
+
+        result[i] = img
+
+    return result.to(torch.float32)
 
 
 def scale_accross_batch_and_channels(tensor, target_size):

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = grad-cam
-version = 1.1.0
+version = 1.4.7
 author = Jacob Gildenblat
 author_email = [email protected]
 description = Many Class Activation Map methods implemented in Pytorch. Including Grad-CAM, Grad-CAM++, Score-CAM, Ablation-CAM and XGrad-CAM
@@ -16,4 +16,4 @@ classifiers =
 
 [options]
 packages = find:
-python_requires = >=3.6
+python_requires = >=3.6

setup.py

@@ -8,7 +8,7 @@
 
 setuptools.setup(
     name='grad-cam',
-    version='1.4.6',
+    version='1.4.7',
     author='Jacob Gildenblat',
     author_email='[email protected]',
     description='Many Class Activation Map methods implemented in Pytorch for classification, segmentation, object detection and more',