Added App files

app.py

@@ -0,0 +1,195 @@
+# Outline
+# Import packages
+# Import modules
+# Constants
+# Load model
+# Function to process user uploaded image/ examples
+# Inference function
+# Gradio examples
+# Gradio App
+
+# Import packages required for the app
+import gradio as gr
+
+# Import custom modules
+import modules.config as config
+import numpy as np
+import torch
+
+# import torchvision
+from modules.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
+
+# Load and initialize the model
+model = CustomResNet()
+
+# Define device
+cpu = torch.device("cpu")
+
+# Using the checkpoint path present in config, load the trained model
+model.load_state_dict(torch.load(config.MODEL_PATH, map_location=cpu), strict=False)
+# Send model to CPU
+model.to(cpu)
+# Make the model in evaluation mode
+model.eval()
+print(f"Model Device: {next(model.parameters()).device}")
+
+
+# Load the misclassified images data
+misclassified_image_data = torch.load(config.MISCLASSIFIED_PATH, map_location=cpu)
+
+# 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 visualization"""
+    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(cpu)
+        print(f"Input Device: {input_image.device}")
+        outputs = model(input_image).to(cpu)
+        print(f"Output Device: {outputs.device}")
+        o = torch.exp(outputs).to(cpu)
+        print(f"Output Exp Device: {o.device}")
+
+        o_np = np.squeeze(np.asarray(o.numpy()))
+        # get indexes of probabilties in descending order
+        sorted_indexes = np.argsort(o_np)[::-1]
+        # sort the probabilities in descending order
+        final_class = classes[o_np.argmax()]
+
+        confidences = {}
+        for cnt in range(int(num_classes)):
+            # set the confidence of highest class with highest probability
+            confidences[classes[sorted_indexes[cnt]]] = float(o_np[sorted_indexes[cnt]])
+
+    # 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, use_cuda=False)
+        grayscale_cam = cam(input_tensor=input_image, targets=None)
+        grayscale_cam = grayscale_cam[0, :]
+        visualization = show_cam_on_image(orginal_img / 255, grayscale_cam, use_rgb=True, image_weight=transparency)
+    else:
+        visualization = orginal_img
+
+    return final_class, confidences, visualization
+
+
+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"""
+
+    # Get the prediction for the input image along with confidence and visualization
+    final_class, confidences, visualization = generate_prediction(
+        input_image, 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 final_class, confidences, visualization, misclassified_fig, gradcam_fig
+
+
+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],
+]
+inference_app = gr.Interface(
+    app_interface,
+    inputs=[
+        # This accepts the image after resizing it to 32x32 which is what our model expects
+        gr.Image(shape=(32, 32)),
+        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, precision=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, precision=0, label="#GradCAM images to show"),
+    ],
+    outputs=[
+        gr.Textbox(label="Top Class", container=True),
+        gr.Label(label="Confidences", container=True),
+        gr.Image(shape=(32, 32), label="Grad CAM/ Input Image", container=True).style(width=256, height=256),
+        gr.Plot(label="Misclassified images", container=True),
+        gr.Plot(label="Grad CAM of Misclassified images"),
+    ],
+    title=TITLE,
+    description=DESCRIPTION,
+    examples=examples,
+)
+inference_app.launch()

assets/model/CustomResNet.pt

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

assets/model/Misclassified_Data.pt

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

gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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

modules/config.py

@@ -0,0 +1,38 @@
+# 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?
+MISCLASSIFIED_PATH = "./assets/model/Misclassified_Data.pt"
+MODEL_PATH = "./assets/model/CustomResNet.pt"
+
+# Set seed value for reproducibility
+SEED = 53
+
+# 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 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",
+    ]
+)
+
+# Needed to load model module
+# What is the start LR and weight decay you'd prefer?
+PREFERRED_START_LR = 5e-3
+PREFERRED_WEIGHT_DECAY = 1e-5

modules/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 12 Model #############
+
+
+# This is for Assignment 12
+# Model used from Assignment 10 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/visualize.py

@@ -0,0 +1,170 @@
+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
+    # Alert: Change the device to cpu for gradio app
+    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=False)
+
+    # 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.4.8
+gradio==3.39.0 
+numpy== 1.25.0
+pillow==9.4.0
+pytorch-lightning==2.0.6
+pytorch==2.0.1
+torch_lr_finder==0.2.1
+torchinfo==1.8.0
+torchmetrics==0.11.4
+torchvision==0.15.2