Added feature for missclassified images

app.py

@@ -1,4 +1,9 @@
-import torch, torchvision
+import torch
+import dataset
+import visualize
+import albumentations
+from utils import get_misclassified_data
+from albumentations.pytorch import ToTensorV2
 from torchvision import transforms
 import numpy as np
 import gradio as gr
@@ -8,12 +13,20 @@ from pytorch_grad_cam.utils.image import show_cam_on_image
 from resnet import ResNet18
 import gradio as gr
 
+cuda = torch.cuda.is_available()
+device = 'cuda' if cuda else 'cpu'
+
 model = ResNet18()
-model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')), strict=False)
+model.load_state_dict(torch.load("model.pth", map_location=torch.device(device)), strict=False)
+
+# dataloader arguments - something you'll fetch these from cmdprmt
+dataloader_args = dict(shuffle=True, batch_size=128, num_workers=4, pin_memory=True) if cuda else dict(shuffle=True, batch_size=64)
+
+test_loader = dataset.get_test_data_loader(**dataloader_args)
 
 inv_normalize = transforms.Normalize(
-    mean=[-0.50/0.23, -0.50/0.23, -0.50/0.23],
-    std=[1/0.23, 1/0.23, 1/0.23]
+    mean=[-0.48215841/0.24348513, -0.44653091/0.26158784, -0.49139968/0.24703223],
+    std=[1/0.24348513, 1/0.26158784, 1/0.24703223]
 )
 classes = ('plane', 'car', 'bird', 'cat', 'deer',
            'dog', 'frog', 'horse', 'ship', 'truck')
@@ -39,14 +52,20 @@ def resize_image_pil(image, new_width, new_height):
 
     return resized
 
-def inference(input_img, is_grad_cam=True, transparency = 0.5, target_layer_number = -1, top_predictions=3):
+def inference(input_img, is_grad_cam=True, transparency = 0.5, target_layer_number = -1,
+              top_predictions=3, is_missclassified_images=True, num_missclassified_images=10):
     input_img = resize_image_pil(input_img, 32, 32)
     
     input_img = np.array(input_img)
     org_img = input_img
     input_img = input_img.reshape((32, 32, 3))
-    transform = transforms.ToTensor()
-    input_img = transform(input_img)
+    transforms = albumentations.Compose(
+        # Normalize
+        [albumentations.Normalize([0.49139968, 0.48215841, 0.44653091], 
+                                  [0.24703223, 0.24348513, 0.26158784]), 
+        # Convert to tensor                          
+        ToTensorV2()])
+    input_img = transforms(image = input_img)['image']
     input_img = input_img
     input_img = input_img.unsqueeze(0)
     outputs = model(input_img)
@@ -68,8 +87,16 @@ def inference(input_img, is_grad_cam=True, transparency = 0.5, target_layer_numb
     
     # Pick the top n predictions
     top_n_confidences = dict(list(sorted_confidences.items())[:top_predictions])
+
+    if is_missclassified_images:
+        # Get the misclassified data from test dataset
+        misclassified_data = get_misclassified_data(model, device, test_loader)
+        # Plot the misclassified data
+        visualize.display_cifar_misclassified_data(misclassified_data, number_of_samples=num_missclassified_images)
+    else:
+        missclassified_images = None
     
-    return classes[prediction[0].item()], visualization, top_n_confidences
+    return classes[prediction[0].item()], visualization, top_n_confidences, missclassified_images
 
 title = "CIFAR10 trained on ResNet18 Model with GradCAM"
 description = "A simple Gradio interface to infer on ResNet model, and get GradCAM results"
@@ -90,12 +117,15 @@ demo = gr.Interface(
         gr.Checkbox(label="Show GradCAM"),
         gr.Slider(0, 1, value = 0.5, label="Overall Opacity of Image"), 
         gr.Slider(-2, -1, value = -2, step=1, label="Which Layer?"),
-        gr.Slider(2, 10, value=3, step=1, label="Number of Top Classes")
+        gr.Slider(2, 10, value=3, step=1, label="Number of Top Classes"),
+        gr.Checkbox(label="Show Misclassified Images"),
+        gr.Slider(5, 40, value=10, step=5, label="Number of Missclassified Images")
         ], 
     outputs = [
         "text", 
         gr.Image(width=256, height=256, label="Output"),
-        gr.Label(label="Top Classes")
+        gr.Label(label="Top Classes"),
+        gr.Plot(label="Missclassified Images")
         ],
     title = title,
     description = description,

dataset.py

@@ -0,0 +1,70 @@
+import numpy as np
+import albumentations
+from torchvision import datasets
+from albumentations.pytorch import ToTensorV2
+from torch.utils.data import Dataset, DataLoader
+
+class CIFAR10Data(Dataset):
+    def __init__(self, dataset, transforms=None) -> None:
+        self.dataset = dataset
+        self.transforms = transforms
+    
+    def __len__(self):
+        return len(self.dataset)
+    
+    def __getitem__(self, index):
+        image, label = self.dataset[index]
+
+        image = np.array(image)
+
+        if self.transforms:
+            image = self.transforms(image=image)['image']
+
+        return image, label
+    
+def _get_test_transforms():
+    test_transforms = albumentations.Compose([albumentations.Normalize([0.49139968, 0.48215841, 0.44653091],
+                                                                     [0.24703223, 0.24348513, 0.26158784]),
+                                            ToTensorV2()])
+    return test_transforms
+
+def _get_data(is_train, is_download):
+    """Method to get data for training or testing
+
+    Args:
+        is_train (bool): True if data is for training else false
+        is_download (bool): True to download dataset from iternet
+
+    Returns:
+        object: Oject of dataset
+    """
+    data = datasets.CIFAR10('../data', train=is_train, download=is_download)
+    return data
+
+def _get_data_loader(data, **kwargs):
+    """Method to get data loader. 
+
+    Args:
+        data (object): Oject of dataset
+
+    Returns:
+        object: Object of DataLoader class used to feed data to neural network model
+    """
+    loader = DataLoader(data, **kwargs)
+    return loader
+
+def get_test_data_loader(**kwargs):
+    """Method to get data loader for testing
+
+    Args:
+        batch_size (int): Number of images in a batch
+
+    Returns:
+        object: Object of DataLoader class used to feed data to neural network model
+    """
+
+    test_transforms = _get_test_transforms()
+    test_data = _get_data(is_train=False, is_download=True)
+    test_data = CIFAR10Data(test_data, test_transforms)
+    test_loader = _get_data_loader(data=test_data, **kwargs)
+    return test_loader

requirements.txt

@@ -1,6 +1,7 @@
 aiofiles==23.2.1
 aiohttp==3.9.5
 aiosignal==1.3.1
+albumentations==1.4.6
 altair==5.3.0
 annotated-types==0.6.0
 anyio==4.3.0
@@ -30,6 +31,7 @@ httpcore==1.0.5
 httpx==0.27.0
 huggingface-hub==0.22.2
 idna==3.7
+imageio==2.34.1
 importlib_resources==6.4.0
 intel-openmp==2021.4.0
 ipykernel==6.29.4
@@ -42,6 +44,7 @@ jsonschema-specifications==2023.12.1
 jupyter_client==8.6.1
 jupyter_core==5.7.2
 kiwisolver==1.4.5
+lazy_loader==0.4
 lightning==2.2.3
 lightning-utilities==0.11.2
 markdown-it-py==3.0.0
@@ -56,6 +59,7 @@ nest-asyncio==1.6.0
 networkx==3.2.1
 numpy==1.26.3
 opencv-python==4.9.0.80
+opencv-python-headless==4.9.0.80
 orjson==3.10.1
 packaging==24.0
 pandas==2.2.2
@@ -74,6 +78,7 @@ python-dateutil==2.9.0.post0
 python-multipart==0.0.9
 pytorch-lightning==2.2.3
 pytz==2024.1
+pywin32==306
 PyYAML==6.0.1
 pyzmq==26.0.2
 referencing==0.35.0
@@ -81,6 +86,7 @@ requests==2.31.0
 rich==13.7.1
 rpds-py==0.18.0
 ruff==0.4.2
+scikit-image==0.23.2
 scikit-learn==1.4.2
 scipy==1.13.0
 semantic-version==2.10.0
@@ -93,14 +99,15 @@ starlette==0.37.2
 sympy==1.12
 tbb==2021.11.0
 threadpoolctl==3.4.0
+tifffile==2024.5.10
 tomlkit==0.12.0
 toolz==0.12.1
-torch==2.3.0
+torch==2.3.0+cu121
 torch-lr-finder==0.2.1
-torchaudio==2.3.0
+torchaudio==2.3.0+cu121
 torchmetrics==1.3.2
 torchsummary==1.5.1
-torchvision==0.18.0
+torchvision==0.18.0+cu121
 tornado==6.4
 tqdm==4.66.2
 traitlets==5.14.3

visualize.py

@@ -0,0 +1,116 @@
+import math
+import numpy as np
+import albumentations
+import matplotlib.pyplot as plt
+
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+
+CLASS_NAMES= ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
+
+def get_inv_transforms():
+    """Method to get transform to inverse the effect of normalization for ploting
+
+    Returns:
+        _Object: Object to apply image augmentations
+    """
+    # Normalize image
+    inv_transforms = albumentations.Normalize([-0.48215841/0.24348513, -0.44653091/0.26158784, -0.49139968/0.24703223],
+                                              [1/0.24348513, 1/0.26158784, 1/0.24703223], max_pixel_value=1.0)
+    return inv_transforms
+
+def plot_samples(train_loader, number_of_images):
+    """Method to plot samples of augmented images
+
+    Args:
+        train_loader (Object): Object of data loader class to get images
+    """
+    inv_transform = get_inv_transforms()
+
+    figure = plt.figure()
+    x_count = 5
+    y_count = 1 if number_of_images <= 5 else math.floor(number_of_images / x_count)
+    images, labels = next(iter(train_loader))
+
+    for index in range(1, number_of_images + 1):
+        plt.subplot(y_count, x_count, index)
+        plt.title(CLASS_NAMES[labels[index].numpy()])
+        plt.axis('off')
+        image = np.array(images[index])
+        image = np.transpose(image, (1, 2, 0))
+        image = inv_transform(image=image)['image']
+        plt.imshow(image)
+
+def display_cifar_misclassified_data(data: list,
+                                     number_of_samples: int = 10):
+    """
+    Function to plot images with labels
+    :param data: List[Tuple(image, label)]
+    :param number_of_samples: Number of images to print
+    """
+    fig = plt.figure(figsize=(10, 10))
+    inv_transform = get_inv_transforms()
+
+    x_count = 5
+    y_count = 1 if number_of_samples <= 5 else math.floor(number_of_samples / x_count)
+
+    for i in range(number_of_samples):
+        plt.subplot(y_count, x_count, i + 1)
+        img = np.array(data[i][0].squeeze().to('cpu'))
+        img = np.transpose(img, (1, 2, 0))
+        img = inv_transform(image=img)['image']
+        plt.imshow(img)
+        plt.title(r"Correct: " + CLASS_NAMES[data[i][1].item()] + '\n' + 'Output: ' + CLASS_NAMES[data[i][2].item()])
+        plt.xticks([])
+        plt.yticks([])
+
+def display_gradcam_output(data: list,
+                           model,
+                           target_layers,
+                           targets=None,
+                           number_of_samples: int = 10,
+                           transparency: float = 0.60):
+    """
+    Function to visualize GradCam output on the data
+    :param data: List[Tuple(image, label)]
+    :param classes: Name of classes in the dataset
+    :param inv_normalize: Mean and Standard deviation values of the dataset
+    :param model: Model architecture
+    :param target_layers: Layers on which GradCam should be executed
+    :param targets: Classes to be focused on for GradCam
+    :param number_of_samples: Number of images to print
+    :param transparency: Weight of Normal image when mixed with activations
+    """
+    # Plot configuration
+    fig = plt.figure(figsize=(10, 10))
+    inv_transform = get_inv_transforms()
+
+    x_count = 5
+    y_count = 1 if number_of_samples <= 5 else math.floor(number_of_samples / x_count)
+
+    # Create an object for GradCam
+    cam = GradCAM(model=model, target_layers=target_layers)
+
+    # Iterate over number of specified images
+    for i in range(number_of_samples):
+        plt.subplot(y_count, x_count, i + 1)
+        input_tensor = data[i][0]
+
+        # Get the activations of the layer for the images
+        grayscale_cam = cam(input_tensor=input_tensor, targets=targets)
+        grayscale_cam = grayscale_cam[0, :]
+
+        # Get back the original image
+        img = np.array(input_tensor.squeeze(0).to('cpu'))
+        img = np.transpose(img, (1, 2, 0))
+        img = inv_transform(image=img)['image']
+        rgb_img = np.clip(img, 0, 1)
+
+        # Mix the activations on the original image
+        visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True, image_weight=transparency)
+
+        # Display the images on the plot
+        plt.imshow(visualization)
+        plt.title(r"Correct: " + CLASS_NAMES[data[i][1].item()] + '\n' + 'Output: ' + CLASS_NAMES[data[i][2].item()])
+        plt.xticks([])
+        plt.yticks([])