Updated resnet class

resnet.py

@@ -24,13 +24,20 @@ import torchvision.transforms as transforms
 import torchvision.datasets as datasets
 import pytorch_lightning as pl
 import matplotlib.pyplot as plt
-
+import matplotlib.gridspec as gridspec
 
 
 PATH_DATASETS = os.environ.get("PATH_DATASETS", ".")
 BATCH_SIZE = 256
 
 
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchsummary import summary
+from io import BytesIO
+import numpy as np
+
 # Model
 class custom_ResNet(pl.LightningModule):
     def __init__(self, data_dir=PATH_DATASETS, learning_rate=2e-4):
@@ -168,7 +175,7 @@ class custom_ResNet(pl.LightningModule):
         acc = pred.eq(y.view_as(pred)).float().mean()
         self.log('test_loss', loss, prog_bar=True)
         self.log('test_acc', acc, prog_bar=True)
-        return pred  # Return predictions instead of loss
+        return pred
 
     def configure_optimizers(self):
         optimizer = torch.optim.Adam(self.parameters(), lr=0.001)
@@ -188,7 +195,7 @@ 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, transform=self.train_transform)
+            cifar_full = CIFAR10(self.data_dir, train=True, download=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)
@@ -215,9 +222,9 @@ class custom_ResNet(pl.LightningModule):
             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
+            misclassified_images.extend(x[~misclassified_mask].detach())
+            misclassified_true_labels.extend(y[~misclassified_mask].detach())
+            misclassified_predicted_labels.extend(pred[~misclassified_mask].detach())
 
             num_collected += sum(~misclassified_mask)
 
@@ -260,15 +267,33 @@ class custom_ResNet(pl.LightningModule):
 
         return gradcam_images
 
-    # Add a 'use_gradcam' parameter to the show_misclassified_images function
+    def create_layout(self, num_images, use_gradcam):
+        num_cols = 3 if use_gradcam else 2
+        fig = plt.figure(figsize=(12, 5 * num_images))
+        gs = gridspec.GridSpec(num_images, num_cols, figure=fig, width_ratios=[0.3, 1, 1] if use_gradcam else [0.5, 1])
+
+        return fig, gs
+
+    def show_images_with_labels(self, fig, gs, i, img, label_text, use_gradcam=False, gradcam_img=None):
+        ax_img = fig.add_subplot(gs[i, 1])
+        ax_img.imshow(img)
+        ax_img.set_title("Original Image")
+        ax_img.axis("off")
+
+        if use_gradcam:
+            ax_gradcam = fig.add_subplot(gs[i, 2])
+            ax_gradcam.imshow(gradcam_img)
+            ax_gradcam.set_title("GradCAM Image")
+            ax_gradcam.axis("off")
+
+        ax_label = fig.add_subplot(gs[i, 0])
+        ax_label.text(0, 0.5, label_text, fontsize=10, verticalalignment='center')
+        ax_label.axis("off")
+
     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
 
-        fig, axs = plt.subplots(num_rows, num_cols, figsize=(8, 5 * num_rows))
+        fig, gs = self.create_layout(num_images, use_gradcam)
 
         if use_gradcam:
             grad_cam_images = self.get_gradcam_images(target_layer=gradcam_layer, transparency=transparency, num_images=num_images)
@@ -277,22 +302,9 @@ class custom_ResNet(pl.LightningModule):
             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
+            # Show true label and predicted label on the left, and images on the right
+            label_text = f"True Label: {self.classes[true_labels[i]]}\nPredicted Label: {self.classes[predicted_labels[i]]}"
+            self.show_images_with_labels(fig, gs, i, img, label_text, use_gradcam, grad_cam_images[i] if use_gradcam else None)
 
+        plt.tight_layout()
+        return fig