Update resnet.py

resnet.py

@@ -2,10 +2,9 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torchsummary import summary
-# imports
+from io import BytesIO
+import numpy as np
 import os
-
-import torch
 from pytorch_lightning import LightningModule, Trainer
 from torch import nn
 from torch.nn import functional as F
@@ -15,7 +14,10 @@ from torchvision import transforms
 from torchvision.datasets import CIFAR10
 from torch_lr_finder import LRFinder
 import math
-
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
+from PIL import Image
 import torch
 from torch.utils.data import DataLoader, random_split
 import torchvision.transforms as transforms
@@ -24,9 +26,11 @@ import pytorch_lightning as pl
 import matplotlib.pyplot as plt
 
 
+
 PATH_DATASETS = os.environ.get("PATH_DATASETS", ".")
 BATCH_SIZE = 256
 
+
 # Model
 class custom_ResNet(pl.LightningModule):
     def __init__(self, data_dir=PATH_DATASETS, learning_rate=2e-4):
@@ -184,12 +188,12 @@ class custom_ResNet(pl.LightningModule):
 
         # Assign train/val datasets for use in dataloaders
         if stage == "fit" or stage is None:
-            cifar_full = CIFAR10(self.data_dir, train=True, download=True, transform=self.train_transform)
+            cifar_full = CIFAR10(self.data_dir, train=True, transform=self.train_transform)
             self.cifar_train, self.cifar_val = random_split(cifar_full, [45000, 5000])
 
         # Assign test dataset for use in dataloader(s)
         if stage == "test" or stage is None:
-            self.cifar_test = CIFAR10(self.data_dir, train=False, download=True, transform=self.test_transform)
+            self.cifar_test = CIFAR10(self.data_dir, train=False, transform=self.test_transform)
 
     def train_dataloader(self):
         return DataLoader(self.cifar_train, batch_size=BATCH_SIZE, num_workers=os.cpu_count())
@@ -208,11 +212,12 @@ class custom_ResNet(pl.LightningModule):
 
         for batch in self.test_dataloader():
             x, y = batch
-            pred = self.forward(x).argmax(dim=1, keepdim=True)
-            misclassified_mask = pred.eq(y.view_as(pred)).squeeze().cpu().numpy()
-            misclassified_images.extend(x[~misclassified_mask])
-            misclassified_true_labels.extend(y[~misclassified_mask])
-            misclassified_predicted_labels.extend(pred[~misclassified_mask])
+            y_hat = self.forward(x)
+            pred = y_hat.argmax(dim=1, keepdim=True)
+            misclassified_mask = pred.eq(y.view_as(pred)).squeeze()
+            misclassified_images.extend(x[~misclassified_mask].detach())  # Detach here to avoid CPU transfer
+            misclassified_true_labels.extend(y[~misclassified_mask].detach())  # Detach here to avoid CPU transfer
+            misclassified_predicted_labels.extend(pred[~misclassified_mask].detach())  # Detach here to avoid CPU transfer
 
             num_collected += sum(~misclassified_mask)
 
@@ -221,36 +226,73 @@ class custom_ResNet(pl.LightningModule):
 
         return misclassified_images[:num_images], misclassified_true_labels[:num_images], misclassified_predicted_labels[:num_images], len(misclassified_images)
 
+
     def normalize_image(self, img_tensor):
         min_val = img_tensor.min()
         max_val = img_tensor.max()
         return (img_tensor - min_val) / (max_val - min_val)
 
-
-
-
-    def show_misclassified_images(self, num_images=10):
+    def get_gradcam_images(self, target_layer=-1, transparency=0.5, num_images=10):
         misclassified_images, true_labels, predicted_labels, num_misclassified = self.collect_misclassified_images(num_images)
+        count = 0
+        k = 0
+        misclassified_images_converted = list()
+        gradcam_images = list()
+
+        if target_layer == -2:
+          target_layer = self.convblock2_l1.cpu()
+        else:
+          target_layer = self.convblock3_l1.cpu()
+
+        dataset_mean, dataset_std = np.array([0.49139968, 0.48215841, 0.44653091]), np.array([0.24703223, 0.24348513, 0.26158784])
+        grad_cam = GradCAM(model=self.cpu(), target_layers=target_layer, use_cuda=False)  # Move model to CPU
+
+        for i in range(0, num_images):
+            img_converted = misclassified_images[i].cpu().numpy().transpose(1, 2, 0)  # Convert tensor to numpy and transpose to (H, W, C)
+            img_converted = dataset_std * img_converted + dataset_mean
+            img_converted = np.clip(img_converted, 0, 1)
+            misclassified_images_converted.append(img_converted)
+            targets = [ClassifierOutputTarget(true_labels[i])]
+            grayscale_cam = grad_cam(input_tensor=misclassified_images[i].unsqueeze(0).cpu(), targets=targets)  # Move input to CPU
+            grayscale_cam = grayscale_cam[0, :]
+            output = show_cam_on_image(img_converted, grayscale_cam, use_rgb=True, image_weight=transparency)
+            gradcam_images.append(output)
+
+        return gradcam_images
+
+    # Add a 'use_gradcam' parameter to the show_misclassified_images function
+    def show_misclassified_images(self, num_images=10, use_gradcam=False, gradcam_layer=-1, transparency=0.5):
+        misclassified_images, true_labels, predicted_labels, num_misclassified = self.collect_misclassified_images(num_images)
+          
+        # Create subplots based on the number of columns required
+        num_rows = num_images
+        num_cols = 2 if use_gradcam else 1  # Show GradCAM images side by side with misclassified images if 'use_gradcam' is True
 
-        num_rows = 2
-        num_cols = math.ceil(num_images / num_rows)
+        fig, axs = plt.subplots(num_rows, num_cols, figsize=(8, 5 * num_rows))
 
-        fig, axs = plt.subplots(num_rows, num_cols, figsize=(5 * num_cols, 5 * num_rows))
-        fig.suptitle(f"Misclassified Images (Showing {num_images} out of {num_misclassified})")
-        plt.subplots_adjust(hspace=0.5)  # Adjust vertical space between subplots
+        if use_gradcam:
+            grad_cam_images = self.get_gradcam_images(target_layer=gradcam_layer, transparency=transparency, num_images=num_images)
 
         for i in range(num_images):
-            img = self.normalize_image(misclassified_images[i]).permute(1, 2, 0)
-            row_idx = i // num_cols
-            col_idx = i % num_cols
-            axs[row_idx, col_idx].imshow(img)
-            axs[row_idx, col_idx].set_title(f"True label: {self.classes[true_labels[i]]}\nPredicted: {self.classes[predicted_labels[i]]}")
-            axs[row_idx, col_idx].axis("off")
-
-        # Remove any empty subplots in the last row (when num_images is not divisible by num_rows)
-        for i in range(num_images, num_rows * num_cols):
-            row_idx = i // num_cols
-            col_idx = i % num_cols
-            axs[row_idx, col_idx].remove()
-
-        plt.show()
+            img = misclassified_images[i].numpy().transpose((1, 2, 0))  # Convert tensor to numpy and transpose to (H, W, C)
+            img = self.normalize_image(img)  # Normalize the image
+
+            if num_cols > 1:  # Use multiple columns for subplots
+                axs[i, 0].imshow(img)
+                axs[i, 0].set_title(f"True label: {self.classes[true_labels[i]]}\nPredicted: {self.classes[predicted_labels[i]]}")
+                axs[i, 0].axis("off")
+
+                if use_gradcam:
+                    # gradcam_img = grad_cam_images[i].numpy().transpose((1, 2, 0))  # Convert tensor to numpy and transpose to (H, W, C)
+                    gradcam_img = self.normalize_image(grad_cam_images[i])  # Normalize the image
+                    axs[i, 1].imshow(gradcam_img)
+                    axs[i, 1].set_title("GradCAM")
+                    axs[i, 1].axis("off")
+            else:  # Use a single column for subplots
+                axs[i].imshow(img)
+                axs[i].set_title(f"True label: {self.classes[true_labels[i]]}\nPredicted: {self.classes[predicted_labels[i]]}")
+                axs[i].axis("off")
+
+        fig.tight_layout()
+        return fig
+