Added files

app.py

@@ -1,10 +1,185 @@
 import gradio as gr
 
-def greet(name):
-    return f"Hello {name}!"
+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
+import modules.config as config
+import numpy as np
+import torch
 
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+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, use_cuda=False)
+        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"""
+
+    # Get the prediction for the input image along with confidence and display_image
+    confidences, display_image = 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 confidences, display_image, misclassified_fig, gradcam_fig
+
+
+
+
+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.Label(label="Confidences", container=True, show_label=True),
+        gr.Image(shape=(32, 32), label="Grad CAM/ Input Image", container=True, show_label=True).style(
+            width=256, height=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()
 
 

modules/config.py

@@ -47,4 +47,8 @@ CIFAR_CLASSES = tuple(
         "ship",
         "truck",
     ]
-)
+)
+
+
+GRADIO_MISCLASSIFIED_PATH = "./assets/model/Misclassified_Data.pt"
+GRADIO_MODEL_PATH = "./assets/model/CustomResNet.pt"