added: model and code and app

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/model/*.pt filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,80 @@
 ---
-title: ERAv2PytorchClassificationLightning
+title: Erav2s13
 emoji: 🔥
-colorFrom: pink
-colorTo: blue
+colorFrom: yellow
+colorTo: red
 sdk: gradio
-sdk_version: 4.31.3
+sdk_version: 4.27.0
 app_file: app.py
 pinned: false
-license: apache-2.0
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Erav2s13- SOUTRIK 🔥
+
+## Overview
+This repository leverages the Hugging Face repository and Gradio for building a user interface (UI). The model training was conducted using Google Colab, and the resulting model files are utilized for inference in the Gradio app.
+
+- **Model Training**: `Main.ipynb` - Colab notebook used to build and train the model.
+- **Inference**: The same model structure and files are used in the Gradio app.
+
+## Custom ResNet Model
+The `custom_resnet.py` file defines a custom ResNet (Residual Network) model using PyTorch Lightning. This model is specifically designed for image classification tasks, particularly for the CIFAR-10 dataset.
+
+### Model Architecture
+The custom ResNet model comprises the following components:
+
+1. **Preparation Layer**: Convolutional layer with 64 filters, followed by batch normalization, ReLU activation, and dropout.
+2. **Layer 1**: Convolutional layer with 128 filters, max pooling, batch normalization, ReLU activation, and dropout. Includes a residual block with two convolutional layers (128 filters each), batch normalization, ReLU activation, and dropout.
+3. **Layer 2**: Convolutional layer with 256 filters, max pooling, batch normalization, ReLU activation, and dropout.
+4. **Layer 3**: Convolutional layer with 512 filters, max pooling, batch normalization, ReLU activation, and dropout. Includes a residual block with two convolutional layers (512 filters each), batch normalization, ReLU activation, and dropout.
+5. **Max Pooling**: Max pooling layer with a kernel size of 4.
+6. **Fully Connected Layer**: Flattened output passed through a fully connected layer with 10 output units (for CIFAR-10 classes).
+7. **Softmax**: Log softmax activation function to obtain predicted class probabilities.
+
+### Training and Evaluation
+The model is trained using PyTorch Lightning, which provides a high-level interface for training, validation, and testing. Key components include:
+
+- **Optimizer**: Adam with a learning rate specified by `PREFERRED_START_LR`.
+- **Scheduler**: OneCycleLR for learning rate adjustment.
+- **Loss and Accuracy**: Cross-entropy loss and accuracy are computed and logged during training, validation, and testing.
+
+### Misclassified Images
+During testing, misclassified images are tracked and stored in a dictionary along with their ground truth and predicted labels, facilitating error analysis and model improvement.
+
+### Hyperparameters
+Key hyperparameters include:
+
+- `PREFERRED_START_LR`: Initial learning rate.
+- `PREFERRED_WEIGHT_DECAY`: Weight decay for regularization.
+
+### Model Summary
+The `detailed_model_summary` function prints a comprehensive summary of the model architecture, detailing input size, kernel size, output size, number of parameters, and trainable status of each layer.
+
+## Lightning Dataset Module
+The `lightning_dataset.py` file contains the `CIFARDataModule` class, which is a PyTorch Lightning `LightningDataModule` for the CIFAR-10 dataset. This class handles data preparation, splitting, and loading.
+
+### CIFARDataModule Class
+
+#### Parameters
+- `data_path`: Directory path for CIFAR-10 dataset.
+- `batch_size`: Batch size for data loaders.
+- `seed`: Random seed for reproducibility.
+- `val_split`: Fraction of training data used for validation (default: 0).
+- `num_workers`: Number of worker processes for data loading (default: 0).
+
+#### Methods
+- `prepare_data`: Downloads CIFAR-10 dataset if not present.
+- `setup`: Defines data transformations and creates training, validation, and testing datasets.
+- `train_dataloader`: Returns training data loader.
+- `val_dataloader`: Returns validation data loader.
+- `test_dataloader`: Returns testing data loader.
+
+#### Utility Methods
+- `_split_train_val`: Splits training dataset into training and validation subsets.
+- `_init_fn`: Initializes random seed for each worker process to ensure reproducibility.
+
+## License
+This project is licensed under the MIT License.
+
+---

app.py

@@ -0,0 +1,208 @@
+import gradio as gr
+
+from models.custom_resnet import CustomResNet
+from modules.visualize import plot_gradcam_images, plot_misclassified_images
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+from torchvision import transforms
+import modules.config as config
+import numpy as np
+import torch
+from PIL import Image
+
+
+TITLE = "CIFAR10 Image classification using a Custom ResNet Model"
+DESCRIPTION = "Gradio App to infer using a Custom ResNet model and get GradCAM results"
+examples = [
+    ["assets/images/airplane.jpg", 3, True, "layer3_x", 0.6, True, 5, True, 5],
+    ["assets/images/bird.jpeg", 4, True, "layer3_x", 0.7, True, 10, True, 20],
+    ["assets/images/car.jpg", 5, True, "layer3_x", 0.5, True, 15, True, 5],
+    ["assets/images/cat.jpeg", 6, True, "layer3_x", 0.65, True, 20, True, 10],
+    ["assets/images/deer.jpg", 7, False, "layer2", 0.75, True, 5, True, 5],
+    ["assets/images/dog.jpg", 8, True, "layer2", 0.55, True, 10, True, 5],
+    ["assets/images/frog.jpeg", 9, True, "layer2", 0.8, True, 15, True, 15],
+    ["assets/images/horse.jpg", 10, False, "layer1_r1", 0.85, True, 20, True, 5],
+    ["assets/images/ship.jpg", 3, True, "layer1_r1", 0.4, True, 5, True, 15],
+    ["assets/images/truck.jpg", 4, True, "layer1_r1", 0.3, True, 5, True, 10],
+]
+
+
+# load and initialise the model
+
+model = CustomResNet()
+
+# Define the device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Using the checkpoint path present in config, load the trained model
+model.load_state_dict(torch.load(config.GRADIO_MODEL_PATH, map_location=device), strict=False)
+# Send model to CPU
+model.to(device)
+# Make the model in evaluation mode
+model.eval()
+
+# Load the misclassified images data
+misclassified_image_data = torch.load(config.GRADIO_MISCLASSIFIED_PATH, map_location=device)
+
+# Class Names
+classes = list(config.CIFAR_CLASSES)
+# Allowed model names
+model_layer_names = ["prep", "layer1_x", "layer1_r1", "layer2", "layer3_x", "layer3_r2"]
+
+
+def get_target_layer(layer_name):
+    """Get target layer for visualization"""
+    if layer_name == "prep":
+        return [model.prep[-1]]
+    elif layer_name == "layer1_x":
+        return [model.layer1_x[-1]]
+    elif layer_name == "layer1_r1":
+        return [model.layer1_r1[-1]]
+    elif layer_name == "layer2":
+        return [model.layer2[-1]]
+    elif layer_name == "layer3_x":
+        return [model.layer3_x[-1]]
+    elif layer_name == "layer3_r2":
+        return [model.layer3_r2[-1]]
+    else:
+        return None
+    
+
+def generate_prediction(input_image, num_classes=3, show_gradcam=True, transparency=0.6, layer_name="layer3_x"):
+    """ "Given an input image, generate the prediction, confidence and display_image"""
+    mean = list(config.CIFAR_MEAN)
+    std = list(config.CIFAR_STD)
+    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean, std)])
+
+    with torch.no_grad():
+        orginal_img = input_image
+        input_image = transform(input_image).unsqueeze(0).to(device)
+        # print(f"Input Device: {input_image.device}")
+        model_output = model(input_image).to(device)
+        # print(f"Output Device: {outputs.device}")
+        output_exp = torch.exp(model_output).to(device)
+        # print(f"Output Exp Device: {o.device}")
+
+        output_numpy = np.squeeze(np.asarray(output_exp.numpy()))
+        # get indexes of probabilties in descending order
+        sorted_indexes = np.argsort(output_numpy)[::-1]
+        # sort the probabilities in descending order
+        # final_class = classes[o_np.argmax()]
+
+        confidences = {}
+        for _ in range(int(num_classes)):
+            # set the confidence of highest class with highest probability
+            confidences[classes[sorted_indexes[_]]] = float(output_numpy[sorted_indexes[_]])
+
+    # Show Grad Cam
+    if show_gradcam:
+        # Get the target layer
+        target_layers = get_target_layer(layer_name)
+        cam = GradCAM(model=model, target_layers=target_layers)
+        cam_generated = cam(input_tensor=input_image, targets=None)
+        cam_generated = cam_generated[0, :]
+        display_image = show_cam_on_image(orginal_img / 255, cam_generated, use_rgb=True, image_weight=transparency)
+
+    else:
+        display_image = orginal_img
+
+    return confidences, display_image
+
+
+def app_interface(
+    input_image,
+    num_classes,
+    show_gradcam,
+    layer_name,
+    transparency,
+    show_misclassified,
+    num_misclassified,
+    show_gradcam_misclassified,
+    num_gradcam_misclassified,
+):
+    """Function which provides the Gradio interface"""
+    input_image = resize_image_pil(input_image, 32, 32)
+
+    input_image = np.array(input_image)
+    org_img = input_image
+    # Get the prediction for the input image along with confidence and display_image
+    confidences, display_image = generate_prediction(org_img, num_classes, show_gradcam, transparency, layer_name)
+
+    if show_misclassified:
+        misclassified_fig, misclassified_axs = plot_misclassified_images(
+            data=misclassified_image_data, class_label=classes, num_images=num_misclassified
+        )
+    else:
+        misclassified_fig = None
+
+    if show_gradcam_misclassified:
+        gradcam_fig, gradcam_axs = plot_gradcam_images(
+            model=model,
+            data=misclassified_image_data,
+            class_label=classes,
+            # Use penultimate block of resnet18 layer 3 as the target layer for gradcam
+            # Decided using model summary so that dimensions > 7x7
+            target_layers=get_target_layer(layer_name),
+            targets=None,
+            num_images=num_gradcam_misclassified,
+            image_weight=transparency,
+        )
+    else:
+        gradcam_fig = None
+
+    # # delete ununsed axises
+    # del misclassified_axs
+    # del gradcam_axs
+
+    return confidences, display_image, misclassified_fig, gradcam_fig
+
+def resize_image_pil(image, new_width, new_height):
+
+    # Convert to PIL image
+    img = Image.fromarray(np.array(image))
+    
+    # Get original size
+    width, height = img.size
+
+    # Calculate scale
+    width_scale = new_width / width
+    height_scale = new_height / height 
+    scale = min(width_scale, height_scale)
+
+    # Resize
+    resized = img.resize((int(width*scale), int(height*scale)), Image.NEAREST)
+    
+    # Crop to exact size
+    resized = resized.crop((0, 0, new_width, new_height))
+
+    return resized
+
+
+
+inference_app = gr.Interface(
+    app_interface,
+    inputs=[
+        # This accepts the image after resizing it to 32x32 which is what our model expects
+        gr.Image(width=256, height=256, label="Input Image"),
+        gr.Number(value=3, maximum=10, minimum=1, step=1.0, precision=0, label="#Classes to show"),
+        gr.Checkbox(True, label="Show GradCAM Image"),
+        gr.Dropdown(model_layer_names, value="layer3_x", label="Visulalization Layer from Model"),
+        # How much should the image be overlayed on the original image
+        gr.Slider(0, 1, 0.6, label="Image Overlay Factor"),
+        gr.Checkbox(True, label="Show Misclassified Images?"),
+        gr.Slider(value=10, maximum=25, minimum=5, step=5.0, label="#Misclassified images to show"),
+        gr.Checkbox(True, label="Visulize GradCAM for Misclassified images?"),
+        gr.Slider(value=10, maximum=25, minimum=5, step=5.0, label="#GradCAM images to show"),
+    ],
+    outputs=[
+        gr.Label(label="Confidences", container=True, show_label=True),
+        gr.Image(label="Grad CAM/ Input Image", container=True, show_label=True,height=256,width=256),
+        gr.Plot(label="Misclassified images", container=True, show_label=True),
+        gr.Plot(label="Grad CAM of Misclassified images", container=True, show_label=True),
+    ],
+    title=TITLE,
+    description=DESCRIPTION,
+    examples=examples,
+)
+inference_app.launch()
+
+

model.py

@@ -0,0 +1,8 @@
+import cv2
+import numpy as np
+
+
+def inverse_pic(input_img):
+    # print(type(input_img))
+    input_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB)
+    return np.flip(input_img)

model/CustomResNet.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe9f24bd3c056c0b4e4c7687678a941ebe0a51ab39ec6d83500ccc02ec2a6574
+size 26326990

model/Misclassified_Data.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:05b440a0bdf7f1996bbca47cd992360b9fb356f195fe7afe38f4ccd047463e58
+size 485025

models/custom_resnet.py

@@ -0,0 +1,456 @@
+"""Module to define the model."""
+
+# Resources
+# https://lightning.ai/docs/pytorch/stable/starter/introduction.html
+# https://lightning.ai/docs/pytorch/stable/starter/converting.html
+# https://lightning.ai/docs/pytorch/stable/notebooks/lightning_examples/cifar10-baseline.html
+
+import modules.config as config
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torchinfo
+from torch.optim.lr_scheduler import OneCycleLR
+from torch_lr_finder import LRFinder
+from torchmetrics import Accuracy
+
+# What is the start LR and weight decay you'd prefer?
+PREFERRED_START_LR = config.PREFERRED_START_LR
+PREFERRED_WEIGHT_DECAY = config.PREFERRED_WEIGHT_DECAY
+
+
+def detailed_model_summary(model, input_size):
+    """Define a function to print the model summary."""
+
+    # https://github.com/TylerYep/torchinfo
+    torchinfo.summary(
+        model,
+        input_size=input_size,
+        batch_dim=0,
+        col_names=(
+            "input_size",
+            "kernel_size",
+            "output_size",
+            "num_params",
+            "trainable",
+        ),
+        verbose=1,
+        col_width=16,
+    )
+
+
+############# Assignment 13 Model #############
+
+
+# This is for Assignment 13
+# Model used from Assignment 11 and converted to lightning model
+class CustomResNet(pl.LightningModule):
+    """This defines the structure of the NN."""
+
+    # Class variable to print shape
+    print_shape = False
+    # Default dropout value
+    dropout_value = 0.02
+
+    def __init__(self):
+        super().__init__()
+
+        # Define loss function
+        # https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html
+        self.loss_function = torch.nn.CrossEntropyLoss()
+
+        # Define accuracy function
+        # https://torchmetrics.readthedocs.io/en/stable/classification/accuracy.html
+        self.accuracy_function = Accuracy(task="multiclass", num_classes=10)
+
+        # Add results dictionary
+        self.results = {
+            "train_loss": [],
+            "train_acc": [],
+            "test_loss": [],
+            "test_acc": [],
+            "val_loss": [],
+            "val_acc": [],
+        }
+
+        # Save misclassified images
+        self.misclassified_image_data = {"images": [], "ground_truths": [], "predicted_vals": []}
+
+        # LR
+        self.learning_rate = PREFERRED_START_LR
+
+        #  Model Notes
+
+        # PrepLayer - Conv 3x3 s1, p1) >> BN >> RELU [64k]
+        # 1. Input size: 32x32x3
+        self.prep = nn.Sequential(
+            nn.Conv2d(
+                in_channels=3,
+                out_channels=64,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+        )
+
+        # Layer1: X = Conv 3x3 (s1, p1) >> MaxPool2D >> BN >> RELU [128k]
+        self.layer1_x = nn.Sequential(
+            nn.Conv2d(
+                in_channels=64,
+                out_channels=128,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+        )
+
+        # Layer1: R1 = ResBlock( (Conv-BN-ReLU-Conv-BN-ReLU))(X) [128k]
+        self.layer1_r1 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=128,
+                out_channels=128,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+            nn.Conv2d(
+                in_channels=128,
+                out_channels=128,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+        )
+
+        # Layer 2: Conv 3x3 [256k], MaxPooling2D, BN, ReLU
+        self.layer2 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=128,
+                out_channels=256,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+        )
+
+        # Layer 3: X = Conv 3x3 (s1, p1) >> MaxPool2D >> BN >> RELU [512k]
+        self.layer3_x = nn.Sequential(
+            nn.Conv2d(
+                in_channels=256,
+                out_channels=512,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.BatchNorm2d(512),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+        )
+
+        # Layer 3: R2 = ResBlock( (Conv-BN-ReLU-Conv-BN-ReLU))(X) [512k]
+        self.layer3_r2 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=512,
+                out_channels=512,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(512),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+            nn.Conv2d(
+                in_channels=512,
+                out_channels=512,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                dilation=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(512),
+            nn.ReLU(),
+            nn.Dropout(self.dropout_value),
+        )
+
+        # MaxPooling with Kernel Size 4
+        # If stride is None, it is set to kernel_size
+        self.maxpool = nn.MaxPool2d(kernel_size=4, stride=4)
+
+        # FC Layer
+        self.fc = nn.Linear(512, 10)
+
+        # Save hyperparameters
+        self.save_hyperparameters()
+
+    def print_view(self, x, msg=""):
+        """Print shape of the model"""
+        if self.print_shape:
+            if msg != "":
+                print(msg, "\n\t", x.shape, "\n")
+            else:
+                print(x.shape)
+
+    def forward(self, x):
+        """Forward pass"""
+
+        # PrepLayer
+        x = self.prep(x)
+        self.print_view(x, "PrepLayer")
+
+        # Layer 1
+        x = self.layer1_x(x)
+        self.print_view(x, "Layer 1, X")
+        r1 = self.layer1_r1(x)
+        self.print_view(r1, "Layer 1, R1")
+        x = x + r1
+        self.print_view(x, "Layer 1, X + R1")
+
+        # Layer 2
+        x = self.layer2(x)
+        self.print_view(x, "Layer 2")
+
+        # Layer 3
+        x = self.layer3_x(x)
+        self.print_view(x, "Layer 3, X")
+        r2 = self.layer3_r2(x)
+        self.print_view(r2, "Layer 3, R2")
+        x = x + r2
+        self.print_view(x, "Layer 3, X + R2")
+
+        # MaxPooling
+        x = self.maxpool(x)
+        self.print_view(x, "Max Pooling")
+
+        # FC Layer
+        # Reshape before FC such that it becomes 1D
+        x = x.view(x.shape[0], -1)
+        self.print_view(x, "Reshape before FC")
+        x = self.fc(x)
+        self.print_view(x, "After FC")
+
+        # Softmax
+        return F.log_softmax(x, dim=-1)
+
+    # Alert: Remove this function later as Tuner is now being used to automatically find the best LR
+    def find_optimal_lr(self, train_loader):
+        """Use LR Finder to find the best starting learning rate"""
+
+        # https://github.com/davidtvs/pytorch-lr-finder
+        # https://github.com/davidtvs/pytorch-lr-finder#notes
+        # https://github.com/davidtvs/pytorch-lr-finder/blob/master/torch_lr_finder/lr_finder.py
+
+        # New optimizer with default LR
+        tmp_optimizer = optim.Adam(self.parameters(), lr=PREFERRED_START_LR, weight_decay=PREFERRED_WEIGHT_DECAY)
+
+        # Create LR finder object
+        lr_finder = LRFinder(self, optimizer=tmp_optimizer, criterion=self.loss_function)
+        lr_finder.range_test(train_loader=train_loader, end_lr=10, num_iter=100)
+        # https://github.com/davidtvs/pytorch-lr-finder/issues/88
+        _, suggested_lr = lr_finder.plot(suggest_lr=True)
+        lr_finder.reset()
+        # plot.figure.savefig("LRFinder - Suggested Max LR.png")
+
+        print(f"Suggested Max LR: {suggested_lr}")
+
+        if suggested_lr is None:
+            suggested_lr = PREFERRED_START_LR
+
+        return suggested_lr
+
+    # optimiser function
+    def configure_optimizers(self):
+        """Add ADAM optimizer to the lightning module"""
+        optimizer = optim.Adam(self.parameters(), lr=self.learning_rate, weight_decay=PREFERRED_WEIGHT_DECAY)
+
+        # Percent start for OneCycleLR
+        # Handles the case where max_epochs is less than 5
+        percent_start = 5 / int(self.trainer.max_epochs)
+        if percent_start >= 1:
+            percent_start = 0.3
+
+        # https://lightning.ai/docs/pytorch/stable/common/optimization.html#total-stepping-batches
+        scheduler_dict = {
+            "scheduler": OneCycleLR(
+                optimizer=optimizer,
+                max_lr=self.learning_rate,
+                total_steps=int(self.trainer.estimated_stepping_batches),
+                pct_start=percent_start,
+                div_factor=100,
+                three_phase=False,
+                anneal_strategy="linear",
+                final_div_factor=100,
+                verbose=False,
+            ),
+            "interval": "step",
+        }
+
+        return {"optimizer": optimizer, "lr_scheduler": scheduler_dict}
+
+    # Define loss function
+    def compute_loss(self, prediction, target):
+        """Compute Loss"""
+
+        # Calculate loss
+        loss = self.loss_function(prediction, target)
+
+        return loss
+
+    # Define accuracy function
+    def compute_accuracy(self, prediction, target):
+        """Compute accuracy"""
+
+        # Calculate accuracy
+        acc = self.accuracy_function(prediction, target)
+
+        return acc * 100
+
+    # Function to compute loss and accuracy for both training and validation
+    def compute_metrics(self, batch):
+        """Function to calculate loss and accuracy"""
+
+        # Get data and target from batch
+        data, target = batch
+
+        # Generate predictions using model
+        pred = self(data)
+
+        # Calculate loss for the batch
+        loss = self.compute_loss(prediction=pred, target=target)
+
+        # Calculate accuracy for the batch
+        acc = self.compute_accuracy(prediction=pred, target=target)
+
+        return loss, acc
+
+    # Get misclassified images based on how many images to return
+    def store_misclassified_images(self):
+        """Get an array of misclassified images"""
+
+        self.misclassified_image_data = {"images": [], "ground_truths": [], "predicted_vals": []}
+
+        # Initialize the model to evaluation mode
+        self.eval()
+
+        # Disable gradient calculation while testing
+        with torch.no_grad():
+            for batch in self.trainer.test_dataloaders:
+                # Move data and labels to device
+                data, target = batch
+                data, target = data.to(self.device), target.to(self.device)
+
+                # Predict using model
+                pred = self(data)
+
+                # Get the index of the max log-probability
+                output = pred.argmax(dim=1)
+
+                # Save the incorrect predictions
+                incorrect_indices = ~output.eq(target)
+
+                # Store images incorrectly predicted, generated predictions and the actual value
+                self.misclassified_image_data["images"].extend(data[incorrect_indices])
+                self.misclassified_image_data["ground_truths"].extend(target[incorrect_indices])
+                self.misclassified_image_data["predicted_vals"].extend(output[incorrect_indices])
+
+    # training function
+    def training_step(self, batch, batch_idx):
+        """Training step"""
+
+        # Compute loss and accuracy
+        loss, acc = self.compute_metrics(batch)
+
+        self.log("train_loss", loss, prog_bar=True, on_epoch=True, logger=True)
+        self.log("train_acc", acc, prog_bar=True, on_epoch=True, logger=True)
+        # Return training loss
+        return loss
+
+    # validation function
+    def validation_step(self, batch, batch_idx):
+        """Validation step"""
+
+        # Compute loss and accuracy
+        loss, acc = self.compute_metrics(batch)
+
+        self.log("val_loss", loss, prog_bar=True, on_epoch=True, logger=True)
+        self.log("val_acc", acc, prog_bar=True, on_epoch=True, logger=True)
+        # Return validation loss
+        return loss
+
+    # test function will just use validation step
+    def test_step(self, batch, batch_idx):
+        """Test step"""
+
+        # Compute loss and accuracy
+        loss, acc = self.compute_metrics(batch)
+
+        self.log("test_loss", loss, prog_bar=False, on_epoch=True, logger=True)
+        self.log("test_acc", acc, prog_bar=False, on_epoch=True, logger=True)
+        # Return validation loss
+        return loss
+
+    # At the end of train epoch append the training loss and accuracy to an instance variable called results
+    def on_train_epoch_end(self):
+        """On train epoch end"""
+
+        # Append training loss and accuracy to results
+        self.results["train_loss"].append(self.trainer.callback_metrics["train_loss"].detach().item())
+        self.results["train_acc"].append(self.trainer.callback_metrics["train_acc"].detach().item())
+
+    # At the end of validation epoch append the validation loss and accuracy to an instance variable called results
+    def on_validation_epoch_end(self):
+        """On validation epoch end"""
+
+        # Append validation loss and accuracy to results
+        self.results["test_loss"].append(self.trainer.callback_metrics["val_loss"].detach().item())
+        self.results["test_acc"].append(self.trainer.callback_metrics["val_acc"].detach().item())
+
+    # # At the end of test epoch append the test loss and accuracy to an instance variable called results
+    # def on_test_epoch_end(self):
+    #     """On test epoch end"""
+
+    #     # Append test loss and accuracy to results
+    #     self.results["test_loss"].append(self.trainer.callback_metrics["test_loss"].detach().item())
+    #     self.results["test_acc"].append(self.trainer.callback_metrics["test_acc"].detach().item())
+
+    # At the end of test save misclassified images, the predictions and ground truth in an instance variable called misclassified_image_data
+    def on_test_end(self):
+        """On test end"""
+
+        print("Test ended! Saving misclassified images")
+        # Get misclassified images
+        self.store_misclassified_images()

modules/config.py

@@ -0,0 +1,54 @@
+# Alert: Change these when running in production
+
+# Constants naming convention: All caps separated by underscore
+# https://realpython.com/python-constants/
+
+# Where do we store the data?
+DATA_PATH = "../../data/"
+CHECKPOINT_PATH = "../../checkpoints/"
+LOGGING_PATH = "../../logs/"
+MISCLASSIFIED_PATH = "../../Misclassified_Data.pt"
+MODEL_PATH = "../../CustomResNet.pt"
+
+# Specify the number of epochs
+NUM_EPOCHS = 24
+
+# Set the batch size
+BATCH_SIZE = 512
+
+# Set seed value for reproducibility
+SEED = 53
+
+# What is the start LR and weight decay you'd prefer?
+PREFERRED_START_LR = 5e-3
+PREFERRED_WEIGHT_DECAY = 1e-5
+
+
+# What is the mean and std deviation of the dataset?
+CIFAR_MEAN = (0.4915, 0.4823, 0.4468)
+CIFAR_STD = (0.2470, 0.2435, 0.2616)
+
+# What is the cutout size?
+CUTOUT_SIZE = 16
+
+# What are the classes in CIFAR10?
+# Create class labels and convert to tuple
+CIFAR_CLASSES = tuple(
+    c.capitalize()
+    for c in [
+        "plane",
+        "car",
+        "bird",
+        "cat",
+        "deer",
+        "dog",
+        "frog",
+        "horse",
+        "ship",
+        "truck",
+    ]
+)
+
+
+GRADIO_MISCLASSIFIED_PATH = "./assets/model/Misclassified_Data.pt"
+GRADIO_MODEL_PATH = "./assets/model/CustomResNet.pt"

modules/dataset.py

@@ -0,0 +1,110 @@
+"""This file contains functions to download and transform the CIFAR10 dataset"""
+# Needed for image transformations
+import albumentations as A
+import modules.config as config
+
+# # Needed for padding issues in albumentations
+# import cv2
+import numpy as np
+from albumentations.pytorch.transforms import ToTensorV2
+from torch.utils.data import Dataset
+
+# Use precomputed values for mean and standard deviation of the dataset
+CIFAR_MEAN = config.CIFAR_MEAN
+CIFAR_STD = config.CIFAR_STD
+CUTOUT_SIZE = config.CUTOUT_SIZE
+
+# Create class labels and convert to tuple
+CIFAR_CLASSES = config.CIFAR_CLASSES
+
+
+class CIFAR10Transforms(Dataset):
+    """Apply albumentations augmentations to CIFAR10 dataset"""
+
+    # Given a dataset and transformations,
+    # apply the transformations and return the dataset
+    def __init__(self, dataset, transforms):
+        self.dataset = dataset
+        self.transforms = transforms
+
+    def __getitem__(self, idx):
+        # Get the image and label from the dataset
+        image, label = self.dataset[idx]
+
+        # Apply transformations on the image
+        image = self.transforms(image=np.array(image))["image"]
+
+        return image, label
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __repr__(self):
+        return f"CIFAR10Transforms(dataset={self.dataset}, transforms={self.transforms})"
+
+    def __str__(self):
+        return f"CIFAR10Transforms(dataset={self.dataset}, transforms={self.transforms})"
+
+
+def apply_cifar_image_transformations(mean=CIFAR_MEAN, std=CIFAR_STD, cutout_size=CUTOUT_SIZE):
+    """
+    Function to apply the required transformations to the MNIST dataset.
+    """
+    # Apply the required transformations to the MNIST dataset
+    train_transforms = A.Compose(
+        [
+            # normalize the images with mean and standard deviation from the whole dataset
+            # https://albumentations.ai/docs/api_reference/augmentations/transforms/#albumentations.augmentations.transforms.Normalize
+            # # transforms.Normalize(cifar_mean, cifar_std),
+            A.Normalize(mean=list(mean), std=list(std)),
+            # RandomCrop 32, 32 (after padding of 4)
+            # https://albumentations.ai/docs/api_reference/augmentations/geometric/transforms/#albumentations.augmentations.geometric.transforms.PadIfNeeded
+            # MinHeight and MinWidth are set to 36 to ensure that the image is padded to 36x36 after padding
+            # border_mode (OpenCV flag): flag that is used to specify the pixel extrapolation method. Should be one of:
+            # cv2.BORDER_CONSTANT, cv2.BORDER_REPLICATE, cv2.BORDER_REFLECT, cv2.BORDER_WRAP, cv2.BORDER_REFLECT_101.
+            # Default: cv2.BORDER_REFLECT_101
+            A.PadIfNeeded(min_height=36, min_width=36),
+            # https://albumentations.ai/docs/api_reference/augmentations/crops/transforms/#albumentations.augmentations.crops.transforms.RandomCrop
+            A.RandomCrop(32, 32),
+            # CutOut(8, 8)
+            # # https://albumentations.ai/docs/api_reference/augmentations/dropout/cutout/#albumentations.augmentations.dropout.cutout.Cutout
+            # # Because we normalized the images with mean and standard deviation from the whole dataset, the fill_value is set to the mean of the dataset
+            # A.Cutout(
+            #     num_holes=1, max_h_size=cutout_size, max_w_size=cutout_size, p=1.0
+            # ),
+            # https://albumentations.ai/docs/api_reference/augmentations/dropout/coarse_dropout/#coarsedropout-augmentation-augmentationsdropoutcoarse_dropout
+            A.CoarseDropout(
+                max_holes=1,
+                max_height=cutout_size,
+                max_width=cutout_size,
+                min_holes=1,
+                min_height=cutout_size,
+                min_width=cutout_size,
+                p=1.0,
+            ),
+            # Convert the images to tensors
+            # # transforms.ToTensor(),
+            ToTensorV2(),
+        ]
+    )
+
+    # Test data transformations
+    test_transforms = A.Compose(
+        # Convert the images to tensors
+        # normalize the images with mean and standard deviation from the whole dataset
+        [
+            A.Normalize(mean=list(mean), std=list(std)),
+            # Convert the images to tensors
+            ToTensorV2(),
+        ]
+    )
+
+    return train_transforms, test_transforms
+
+
+def calculate_mean_std(dataset):
+    """Function to calculate the mean and standard deviation of CIFAR dataset"""
+    data = dataset.data.astype(np.float32) / 255.0
+    mean = np.mean(data, axis=(0, 1, 2))
+    std = np.std(data, axis=(0, 1, 2))
+    return mean, std

modules/lightning_dataset.py

@@ -0,0 +1,109 @@
+"""This file contains functions to prepare dataloader in the way lightning expects"""
+import pytorch_lightning as pl
+import torchvision.datasets as datasets
+from lightning_fabric.utilities.seed import seed_everything
+from modules.dataset import CIFAR10Transforms, apply_cifar_image_transformations
+from torch.utils.data import DataLoader, random_split
+
+
+class CIFARDataModule(pl.LightningDataModule):
+    """Lightning DataModule for CIFAR10 dataset"""
+
+    def __init__(self, data_path, batch_size, seed, val_split=0, num_workers=0):
+        super().__init__()
+
+        self.data_path = data_path
+        self.batch_size = batch_size
+        self.seed = seed
+        self.val_split = val_split
+        self.num_workers = num_workers
+        self.dataloader_dict = {
+            # "shuffle": True,
+            "batch_size": self.batch_size,
+            "num_workers": self.num_workers,
+            "pin_memory": True,
+            # "worker_init_fn": self._init_fn,
+            "persistent_workers": self.num_workers > 0,
+        }
+        self.prepare_data_per_node = False
+
+        # Fixes attribute defined outside __init__ warning
+        self.training_dataset = None
+        self.validation_dataset = None
+        self.testing_dataset = None
+
+        # # Make sure data is downloaded
+        # self.prepare_data()
+
+    def _split_train_val(self, dataset):
+        """Split the dataset into train and validation sets"""
+
+        # Throw an error if the validation split is not between 0 and 1
+        if not 0 < self.val_split < 1:
+            raise ValueError("Validation split must be between 0 and 1")
+
+        # # Set seed again, might not be necessary
+        # seed_everything(int(self.seed))
+
+        # Calculate lengths of each dataset
+        total_length = len(dataset)
+        train_length = int((1 - self.val_split) * total_length)
+        val_length = total_length - train_length
+
+        # Split the dataset
+        train_dataset, val_dataset = random_split(dataset, [train_length, val_length])
+
+        return train_dataset, val_dataset
+
+    # https://lightning.ai/docs/pytorch/stable/data/datamodule.html#prepare-data
+    def prepare_data(self):
+        # Download the CIFAR10 dataset if it doesn't exist
+        datasets.CIFAR10(self.data_path, train=True, download=True)
+        datasets.CIFAR10(self.data_path, train=False, download=True)
+
+    # https://lightning.ai/docs/pytorch/stable/data/datamodule.html#setup
+    # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.core.hooks.DataHooks.html#lightning.pytorch.core.hooks.DataHooks.setup
+    def setup(self, stage=None):
+        # seed_everything(int(self.seed))
+
+        # Define the data transformations
+        train_transforms, test_transforms = apply_cifar_image_transformations()
+        val_transforms = test_transforms
+
+        # Create train and validation datasets
+        if stage == "fit" or stage is None:
+            if self.val_split != 0:
+                # Split the training data into training and validation sets
+                data_train, data_val = self._split_train_val(datasets.CIFAR10(self.data_path, train=True))
+                # Apply transformations
+                self.training_dataset = CIFAR10Transforms(data_train, train_transforms)
+                self.validation_dataset = CIFAR10Transforms(data_val, val_transforms)
+            else:
+                # Only training data here
+                self.training_dataset = CIFAR10Transforms(
+                    datasets.CIFAR10(self.data_path, train=True), train_transforms
+                )
+                # Validation will be same sa test
+                self.validation_dataset = CIFAR10Transforms(
+                    datasets.CIFAR10(self.data_path, train=False), val_transforms
+                )
+
+        # Create test dataset
+        if stage == "test" or stage is None:
+            # Assign Test split(s) for use in Dataloaders
+            self.testing_dataset = CIFAR10Transforms(datasets.CIFAR10(self.data_path, train=False), test_transforms)
+
+    # https://lightning.ai/docs/pytorch/stable/data/datamodule.html#train-dataloader
+    def train_dataloader(self):
+        return DataLoader(self.training_dataset, **self.dataloader_dict, shuffle=True)
+
+    # https://lightning.ai/docs/pytorch/stable/data/datamodule.html#val-dataloader
+    def val_dataloader(self):
+        return DataLoader(self.validation_dataset, **self.dataloader_dict, shuffle=False)
+
+    # https://lightning.ai/docs/pytorch/stable/data/datamodule.html#test-dataloader
+    def test_dataloader(self):
+        return DataLoader(self.testing_dataset, **self.dataloader_dict, shuffle=False)
+
+    def _init_fn(self, worker_id):
+        seed_everything(int(self.seed) + worker_id)

modules/trainer.py

@@ -0,0 +1,121 @@
+"""Module to define the train and test functions."""
+
+# from functools import partial
+
+import modules.config as config
+import pytorch_lightning as pl
+import torch
+from modules.utils import create_folder_if_not_exists
+from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint, ModelSummary
+
+# Import tuner
+from pytorch_lightning.tuner.tuning import Tuner
+
+# What is the start LR and weight decay you'd prefer?
+PREFERRED_START_LR = config.PREFERRED_START_LR
+
+
+def train_and_test_model(
+    batch_size,
+    num_epochs,
+    model,
+    datamodule,
+    logger,
+    debug=False,
+):
+    """Trains and tests the model by iterating through epochs using Lightning Trainer."""
+
+    print(f"\n\nBatch size: {batch_size}, Total epochs: {num_epochs}\n\n")
+
+    print("Defining Lightning Callbacks")
+
+    # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html#modelcheckpoint
+    checkpoint = ModelCheckpoint(
+        dirpath=config.CHECKPOINT_PATH, monitor="val_acc", mode="max", filename="model_best_epoch", save_last=True
+    )
+    # # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.LearningRateMonitor.html#learningratemonitor
+    lr_rate_monitor = LearningRateMonitor(logging_interval="epoch", log_momentum=False)
+    # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelSummary.html#lightning.pytorch.callbacks.ModelSummary
+    model_summary = ModelSummary(max_depth=0)
+
+    print("Defining Lightning Trainer")
+    # Change trainer settings for debugging
+    if debug:
+        num_epochs = 1
+        fast_dev_run = True
+        overfit_batches = 0.1
+        profiler = "advanced"
+    else:
+        fast_dev_run = False
+        overfit_batches = 0.0
+        profiler = None
+
+    # https://lightning.ai/docs/pytorch/stable/common/trainer.html#methods
+    trainer = pl.Trainer(
+        precision=16,
+        fast_dev_run=fast_dev_run,
+        # deterministic=True,
+        # devices="auto",
+        # accelerator="auto",
+        max_epochs=num_epochs,
+        logger=logger,
+        # enable_model_summary=False,
+        overfit_batches=overfit_batches,
+        log_every_n_steps=10,
+        # num_sanity_val_steps=5,
+        profiler=profiler,
+        # check_val_every_n_epoch=1,
+        callbacks=[checkpoint, lr_rate_monitor, model_summary],
+        # callbacks=[checkpoint],
+    )
+
+    # # Using the learning rate finder
+    # model.learning_rate = model.find_optimal_lr(train_loader=datamodule.train_dataloader())
+
+    # Using the lr_find from Trainer.tune method instead
+    # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.tuner.tuning.Tuner.html#lightning.pytorch.tuner.tuning.Tuner
+    # https://www.youtube.com/watch?v=cLZv0eZQSIE
+    print("Finding the optimal learning rate using Lightning Tuner.")
+    tuner = Tuner(trainer)
+    tuner.lr_find(
+        model=model,
+        datamodule=datamodule,
+        min_lr=PREFERRED_START_LR,
+        max_lr=5,
+        num_training=200,
+        mode="linear",
+        early_stop_threshold=10,
+        attr_name="learning_rate",
+    )
+
+    trainer.fit(model, datamodule=datamodule)
+    trainer.test(model, dataloaders=datamodule.test_dataloader())
+
+    # # Obtain the results dictionary from model
+    print("Collecting epoch level model results.")
+    results = model.results
+    # print(f"Results Length: {len(results)}")
+
+    # Get the list of misclassified images
+    print("Collecting misclassified images.")
+    misclassified_image_data = model.misclassified_image_data
+    # print(f"Misclassified Images Length: {len(misclassified_image_data)}")
+
+    # Save the model using torch save as backup
+    print("Saving the model.")
+    print(f"Model saved to {config.MODEL_PATH}")
+    create_folder_if_not_exists(config.MODEL_PATH)
+    torch.save(model.state_dict(), config.MODEL_PATH)
+
+    # Save first few misclassified images data to a file
+    num_elements = 20
+    print(f"Saving first {num_elements} misclassified images.")
+    subset_misclassified_image_data = {"images": [], "ground_truths": [], "predicted_vals": []}
+    subset_misclassified_image_data["images"] = misclassified_image_data["images"][:num_elements]
+    subset_misclassified_image_data["ground_truths"] = misclassified_image_data["ground_truths"][:num_elements]
+    subset_misclassified_image_data["predicted_vals"] = misclassified_image_data["predicted_vals"][:num_elements]
+    create_folder_if_not_exists(config.MISCLASSIFIED_PATH)
+    torch.save(subset_misclassified_image_data, config.MISCLASSIFIED_PATH)
+
+    return trainer, results, misclassified_image_data
+    # return trainer

modules/utils.py

@@ -0,0 +1,71 @@
+"""Module to define utility functions for the project."""
+import os
+
+import torch
+
+
+def get_num_workers(model_run_location):
+    """Given a run mode, return the number of workers to be used for data loading."""
+
+    # calculate the number of workers
+    num_workers = (os.cpu_count() - 1) if os.cpu_count() > 3 else 2
+
+    # If run_mode is local, use only 2 workers
+    num_workers = num_workers if model_run_location == "colab" else 0
+
+    return num_workers
+
+
+# Function to save the model
+# https://debuggercafe.com/saving-and-loading-the-best-model-in-pytorch/
+def save_model(epoch, model, optimizer, scheduler, batch_size, criterion, file_name):
+    """
+    Function to save the trained model along with other information to disk.
+    """
+    # print(f"Saving model from epoch {epoch}...")
+    torch.save(
+        {
+            "epoch": epoch,
+            "model_state_dict": model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "scheduler_state_dict": scheduler.state_dict(),
+            "batch_size": batch_size,
+            "loss": criterion,
+        },
+        file_name,
+    )
+
+
+# Given a list of train_losses, train_accuracies, test_losses,
+# test_accuracies, loop through epoch and print the metrics
+def pretty_print_metrics(num_epochs, results):
+    """
+    Function to print the metrics in a pretty format.
+    """
+    # Extract train_losses, train_acc, test_losses, test_acc from results
+    train_losses = results["train_loss"]
+    train_acc = results["train_acc"]
+    test_losses = results["test_loss"]
+    test_acc = results["test_acc"]
+
+    for i in range(num_epochs):
+        print(
+            f"Epoch: {i+1:02d}, Train Loss: {train_losses[i]:.4f}, "
+            f"Test Loss: {test_losses[i]:.4f}, Train Accuracy: {train_acc[i]:.4f}, "
+            f"Test Accuracy: {test_acc[i]:.4f}"
+        )
+
+
+# Given a file path, extract the folder path and create folder recursively if it does not already exist
+def create_folder_if_not_exists(file_path):
+    """
+    Function to create a folder if it does not exist.
+    """
+    # Extract the folder path
+    folder_path = os.path.dirname(file_path)
+    print(f"Folder path: {folder_path}")
+
+    # Create the folder if it does not exist
+    if not os.path.exists(folder_path):
+        os.makedirs(folder_path,exist_ok=True)
+        print(f"Created folder: {folder_path}")

modules/visualize.py

@@ -0,0 +1,169 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+
+
+def convert_back_image(image):
+    """Using mean and std deviation convert image back to normal"""
+    cifar10_mean = (0.4914, 0.4822, 0.4471)
+    cifar10_std = (0.2469, 0.2433, 0.2615)
+    image = image.numpy().astype(dtype=np.float32)
+
+    for i in range(image.shape[0]):
+        image[i] = (image[i] * cifar10_std[i]) + cifar10_mean[i]
+
+    # To stop throwing a warning that image pixels exceeds bounds
+    image = image.clip(0, 1)
+
+    return np.transpose(image, (1, 2, 0))
+
+
+def plot_sample_training_images(batch_data, batch_label, class_label, num_images=30):
+    """Function to plot sample images from the training data."""
+    images, labels = batch_data, batch_label
+
+    # Calculate the number of images to plot
+    num_images = min(num_images, len(images))
+    # calculate the number of rows and columns to plot
+    num_cols = 5
+    num_rows = int(np.ceil(num_images / num_cols))
+
+    # Initialize a subplot with the required number of rows and columns
+    fig, axs = plt.subplots(num_rows, num_cols, figsize=(10, 10))
+
+    # Iterate through the images and plot them in the grid along with class labels
+
+    for img_index in range(1, num_images + 1):
+        plt.subplot(num_rows, num_cols, img_index)
+        plt.tight_layout()
+        plt.axis("off")
+        plt.imshow(convert_back_image(images[img_index - 1]))
+        plt.title(class_label[labels[img_index - 1].item()])
+        plt.xticks([])
+        plt.yticks([])
+
+    return fig, axs
+
+
+def plot_train_test_metrics(results):
+    """
+    Function to plot the training and test metrics.
+    """
+    # Extract train_losses, train_acc, test_losses, test_acc from results
+    train_losses = results["train_loss"]
+    train_acc = results["train_acc"]
+    test_losses = results["test_loss"]
+    test_acc = results["test_acc"]
+
+    # Plot the graphs in a 1x2 grid showing the training and test metrics
+    fig, axs = plt.subplots(1, 2, figsize=(16, 8))
+
+    # Loss plot
+    axs[0].plot(train_losses, label="Train")
+    axs[0].plot(test_losses, label="Test")
+    axs[0].set_title("Loss")
+    axs[0].legend(loc="upper right")
+
+    # Accuracy plot
+    axs[1].plot(train_acc, label="Train")
+    axs[1].plot(test_acc, label="Test")
+    axs[1].set_title("Accuracy")
+    axs[1].legend(loc="upper right")
+
+    return fig, axs
+
+
+def plot_misclassified_images(data, class_label, num_images=10):
+    """Plot the misclassified images from the test dataset."""
+    # Calculate the number of images to plot
+    num_images = min(num_images, len(data["ground_truths"]))
+    # calculate the number of rows and columns to plot
+    num_cols = 5
+    num_rows = int(np.ceil(num_images / num_cols))
+
+    # Initialize a subplot with the required number of rows and columns
+    fig, axs = plt.subplots(num_rows, num_cols, figsize=(num_cols * 2, num_rows * 2))
+
+    # Iterate through the images and plot them in the grid along with class labels
+
+    for img_index in range(1, num_images + 1):
+        # Get the ground truth and predicted labels for the image
+        label = data["ground_truths"][img_index - 1].cpu().item()
+        pred = data["predicted_vals"][img_index - 1].cpu().item()
+        # Get the image
+        image = data["images"][img_index - 1].cpu()
+        # Plot the image
+        plt.subplot(num_rows, num_cols, img_index)
+        plt.tight_layout()
+        plt.axis("off")
+        plt.imshow(convert_back_image(image))
+        plt.title(f"""ACT: {class_label[label]} \nPRED: {class_label[pred]}""")
+        plt.xticks([])
+        plt.yticks([])
+
+    return fig, axs
+
+
+# Function to plot gradcam for misclassified images using pytorch_grad_cam
+def plot_gradcam_images(
+    model,
+    data,
+    class_label,
+    target_layers,
+    targets=None,
+    num_images=10,
+    image_weight=0.25,
+):
+    """Show gradcam for misclassified images"""
+
+    # Calculate the number of images to plot
+    num_images = min(num_images, len(data["ground_truths"]))
+    # calculate the number of rows and columns to plot
+    num_cols = 5
+    num_rows = int(np.ceil(num_images / num_cols))
+
+    # Initialize a subplot with the required number of rows and columns
+    fig, axs = plt.subplots(num_rows, num_cols, figsize=(num_cols * 2, num_rows * 2))
+
+    # Initialize the GradCAM object
+    # https://github.com/jacobgil/pytorch-grad-cam/blob/master/pytorch_grad_cam/grad_cam.py
+    # https://github.com/jacobgil/pytorch-grad-cam/blob/master/pytorch_grad_cam/base_cam.py
+    cam = GradCAM(model=model, target_layers=target_layers)
+
+    # Iterate through the images and plot them in the grid along with class labels
+    for img_index in range(1, num_images + 1):
+        # Extract elements from the data dictionary
+        # Get the ground truth and predicted labels for the image
+        label = data["ground_truths"][img_index - 1].cpu().item()
+        pred = data["predicted_vals"][img_index - 1].cpu().item()
+        # Get the image
+        image = data["images"][img_index - 1].cpu()
+
+        # Get the GradCAM output
+        # https://github.com/jacobgil/pytorch-grad-cam/blob/master/pytorch_grad_cam/utils/model_targets.py
+        grad_cam_output = cam(
+            input_tensor=image.unsqueeze(0),
+            targets=targets,
+            aug_smooth=True,
+            eigen_smooth=True,
+        )
+        grad_cam_output = grad_cam_output[0, :]
+
+        # Overlay gradcam on top of numpy image
+        overlayed_image = show_cam_on_image(
+            convert_back_image(image),
+            grad_cam_output,
+            use_rgb=True,
+            image_weight=image_weight,
+        )
+
+        # Plot the image
+        plt.subplot(num_rows, num_cols, img_index)
+        plt.tight_layout()
+        plt.axis("off")
+        plt.imshow(overlayed_image)
+        plt.title(f"""ACT: {class_label[label]} \nPRED: {class_label[pred]}""")
+        plt.xticks([])
+        plt.yticks([])
+    return fig, axs

requirements.txt

@@ -0,0 +1,11 @@
+albumentations==1.3.1
+grad-cam==1.5.0
+gradio==3.39.0 
+numpy== 1.25.0
+pillow==9.4.0
+pytorch-lightning==2.0.6
+torch_lr_finder==0.2.1
+torch==2.0.1
+torchinfo==1.8.0
+torchmetrics==0.11.4
+torchvision==0.15.2