app.py

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()