app.py

import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
import torch
import random
from collections import OrderedDict
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
import numpy as np
from PIL import Image
from torchvision import transforms

dropout_value = 0.1
class ResBlock(nn.Module):
  def __init__(self, in_channels, out_channels):
    super(ResBlock,self).__init__()
    self.res_block = nn.Sequential(
        nn.Conv2d(in_channels=in_channels, out_channels = out_channels, kernel_size=3, stride =1 , padding =1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(),
        nn.Conv2d(in_channels=out_channels, out_channels = out_channels, kernel_size=3, stride =1 , padding =1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(),
    )

  def forward (self, x):
    x = self.res_block(x)
    return x


class LayerBlock(nn.Module):
  def __init__(self, in_channels, out_channels):
    super(LayerBlock,self).__init__()
    self.layer_block = nn.Sequential(
        nn.Conv2d(in_channels=in_channels, out_channels = out_channels, kernel_size=3, stride =1 , padding =1),
        nn.MaxPool2d(kernel_size=2,stride=2),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(),
    )

  def forward (self, x):
    x = self.layer_block(x)
    return x

class custom_resnet_s10(nn.Module):
  def __init__(self, num_classes=10):
    super(custom_resnet_s10,self).__init__()

    self.PrepLayer = nn.Sequential(
        nn.Conv2d(in_channels = 3, out_channels=64, kernel_size = 3, stride = 1, padding =1),
        nn.BatchNorm2d(64),
        nn.ReLU(),
    )
    self.Layer1 = LayerBlock(in_channels = 64, out_channels=128)
    self.resblock1 = ResBlock(in_channels =128, out_channels=128)
    self.Layer2 = LayerBlock(in_channels = 128, out_channels=256)
    self.resblock2 = ResBlock(in_channels =256, out_channels=256)
    self.Layer3 = LayerBlock(in_channels = 256, out_channels=512)
    self.resblock3 = ResBlock(in_channels =512, out_channels=512)
    self.max_pool4 = nn.MaxPool2d(kernel_size=4, stride=4) # 512,512, 4/4 = 512,512,1
    self.fc = nn.Linear(512,num_classes)

  def forward(self,x):
    x = self.PrepLayer(x)
    x = self.Layer1(x)
    resl1 = self.resblock1(x)
    
    x = x+resl1
    x = self.Layer2(x)
    resl2 = self.resblock2(x)
    x = x+resl2
    
    
    x = self.Layer3(x)
    
    resl3 = self.resblock3(x)
    x = x+resl3
    x = self.max_pool4(x)
    x = x.view(x.size(0),-1)
    
    x = self.fc(x)
    return x

def get_device():
    if torch.cuda.is_available():
        device = "cuda"
    elif torch.backends.mps.is_available():
        device = "mps"
    else:
        device = "cpu"
    print("Device Selected:", device)
    return device

DEVICE = get_device()

# Load the list of tensors from the file
loaded_misclassified_image_list = torch.load('misclassified_images_list.pt')

# Instantiate the model (make sure it has the same architecture)
loaded_model = custom_resnet_s10()
loaded_model = loaded_model.to(DEVICE)

# Load the saved state dictionary
loaded_model.load_state_dict(torch.load('model.pth', map_location=DEVICE), strict=False)

# Put the loaded model in evaluation mode
loaded_model.eval()

classes = ['plane', 'car', 'bird', 'cat', 'deer','dog', 'frog', 'horse', 'ship', 'truck']
mean = (0.49139968, 0.48215827, 0.44653124)
std = (0.24703233, 0.24348505, 0.26158768)
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=mean, std=std)
])

dict_layer = {'layer3': loaded_model.resblock2.res_block[-1],
                'layer4': loaded_model.resblock3.res_block[-1]}

def view_gradcam_images(choice_gradcam):
    if choice_gradcam == "Yes (View Existing Images)":
      return gr.update(label ="Number of GradCAM Images to view", visible=True, interactive = True), \
      gr.update(visible=True), \
      gr.update(visible=True),  gr.update(visible=True), \
      gr.update(visible=False) # Gallery not shown as yet
    else:
      #TODO: to be completed
      return gr.update(visible=False), gr.update(visible=False),gr.update(visible=False),gr.update(visible=False),gr.update(visible=False)

def process_gradcam_images(num_images,layer,opacity,image_list=None):
 
  if not image_list:
    selected_data = random.sample(loaded_misclassified_image_list, min(num_images,len(loaded_misclassified_image_list)))
  else:
    selected_data = [image_list]

  layer_model = dict_layer.get(layer)
  cam = GradCAM(model=loaded_model, target_layers = [layer_model], use_cuda = False)
  grad_images = []
  inv_normalize = transforms.Normalize(
    mean=[-0.50/0.2197, -0.50/0.1858, -0.50/0.1569], # mean_ds = [0.2197, 0.1858, 0.1569]
    std=[1/0.1810, 1/0.1635, 1/0.1511] # std_dev_ds =[0.1810, 0.1635, 0.1511]
    )
  for i, (img, pred, correct) in enumerate(selected_data):
    input_tensor = img.unsqueeze(0)
    targets = [ClassifierOutputTarget(pred)]
    grayscale_cam = cam(input_tensor=input_tensor, targets=targets)

    grayscale_cam = grayscale_cam[0, :]

    # Get back the original image
    img = input_tensor.squeeze(0).to('cpu')
    img = inv_normalize(img)
    rgb_img = np.transpose(img, (1, 2, 0))
    rgb_img = torch.clamp(rgb_img, max = 1)
    rgb_img = rgb_img.numpy()

    visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True, image_weight=opacity)
    if not image_list:
        grad_images.append(((visualization),f'Pred: {classes[pred.cpu()]} | Truth :{classes[correct.cpu()]}'))
    else:
        grad_images.append(((visualization),f'Prediction: {classes[pred.cpu()]}'))

  print(str(num_images) + "**" + str(layer) + "**" + str(opacity))
  return  grad_images, gr.update(visible=True)



def process_misclassified_images(num_images):
  selected_data = random.sample(loaded_misclassified_image_list, min(num_images,len(loaded_misclassified_image_list)))
  misclassified_images = []
  for i, (img, pred, correct) in enumerate(selected_data):
    img, pred, target = img.cpu().numpy().astype(dtype=np.float32), pred.cpu(), correct.cpu()
    for j in range(img.shape[0]):
        img[j] = (img[j] * std[j]) + mean[j]
    img = np.transpose(img, (1, 2, 0))
    img = Image.fromarray((img * 255).astype(np.uint8))
    misclassified_images.append(((img),f'Pred: {classes[pred]} | Truth :{classes[correct]}'))
  return  misclassified_images, gr.update(visible=True)

def view_misclassified_images(choice_misclassified):
    if choice_misclassified == "Yes":
      return gr.update(label ="Number of Misclassified Images to view", visible=True, interactive = True),gr.update(visible=True),gr.update(visible=False)
    else:
        return gr.update(visible=False),gr.update(visible=False),gr.update(visible=False)

def classify_image(image, num_classes=3, grad_cam_choice = False, layer = None, opacity = 0.8 ):
  # transforming image and getting prediction from model
  transformed_image = transform(image)
  image_tensor = transformed_image.to(DEVICE).unsqueeze(0)#transform(torch.tensor(image).to(DEVICE)).unsqueeze(0) # making it a batch

  # sending it to model to get prediction
  logits = loaded_model(image_tensor) # logits
  output = F.softmax(logits.view(-1)) #F.softmax(output.flatten(), dim=-1) #

  confidences = [(classes[i], float(output[i])) for i in range(len(classes))]
  confidences.sort(key=lambda x: x[1], reverse=True)
  confidences = OrderedDict(confidences[:num_classes])
  label = torch.argmax(output).item()

  if grad_cam_choice:
    print("** Before Calling **",transformed_image.shape)
    image_list = [transformed_image.to(DEVICE),torch.tensor(label).to(DEVICE),torch.tensor(label).to(DEVICE)]

    grad_cam_output,_ = process_gradcam_images(num_images = 1,layer = layer,opacity= opacity,image_list=image_list)

    return confidences, grad_cam_output , gr.update(visible=True)
  else:
    return confidences, gr.update(visible=False),gr.update(visible=False)


with gr.Blocks() as demo:
  with gr.Tab("GradCam"):
    gr.Markdown(
                """
                Visualize Class Activations Maps (helps to see what the model is actually looking at in the image) generated by the model's layer for the predicted class
                - For existing images
                - For new images (choose an example image or upload your own)
                """
            )
    with gr.Column():
      with gr.Box():
        radio_gradcam = gr.Radio(["Yes (View Existing Images)", "No (New or Example Images)"], label="Do you want to view existing GradCAM images?")
      with gr.Column():
        with gr.Row():
          slider_gradcam_num_images = gr.Slider(minimum=1, maximum =10, value = 1, step =1, visible= False, interactive = False)
          dropdown_gradcam_layer = gr.Dropdown(choices=['layer4', 'layer3'], value = "layer4", label="Please select the layer from which the GradCAM would be taken", interactive = True, visible= False)
          slider_gradcam_opacity = gr.Slider(label ="Opacity of Images", minimum=0.05, maximum =1.00, value = 0.70, step =0.05, visible= False, interactive = True)
        button_gradcam = gr.Button("View GradCAM Output", visible = False)
        # txt_gradcam = gr.Textbox ("GradCAM output here" , visible = True)
        output_gallery_gradcam=gr.Gallery(label="GradCAM Output", min_width=512,columns=4, visible = False)
      with gr.Box():
        with gr.Row():
          with gr.Column():
            input_image_classify = gr.Image(label="Classification",type="pil", shape=(32, 32))
            slider_classify_num_classes = gr.Slider(label="Select the number of top classes to be shown",minimum=1, maximum =10, value = 3, step = 1, visible= True, interactive = True)
            checkbox_gradcam_classify = gr.Checkbox(label="Enable GradCAM", value=True, info="Do you want to see Class Activation Maps?", visible=True)
            # txt_classify= gr.Textbox ("Classification output here" , visible = True)
            dropdown_gradcam_classify_layer = gr.Dropdown(choices=['layer4', 'layer3'], value = "layer4", label="Please select the layer from which the GradCAM would be taken", interactive = True, visible= True)
            slider_gradcam_classify_opacity = gr.Slider(label ="Opacity of Images", minimum=0.05, maximum =1.00, value = 0.80, step =0.05, visible= True, interactive = True)
            button_classify = gr.Button("Submit to Classify Image", visible = True)

          with gr.Column():
            label_classify = gr.Label(num_top_classes=10, visible = True)
            gallery_gradcam_classify = gr.Gallery(label="GradCAM Output", min_width=256,columns=1, visible = True)
        with gr.Row():
          gr.Examples(['bird1.jpg','car1.jpg','deer1.jpg','frog1.jpg','plane1.jpg','ship1.jpg','truck1.jpg',"cat1.jpg","dog1.jpg","horse1.jpg"],inputs=[input_image_classify])
  with gr.Tab("Misclassified Examples"):
    gr.Markdown(
        """
        The AI model is not able to predict correct image labels all the time.

        Select "Yes" to visualize the misclassified images with their model predicted label and ground truth label.
        """
    )
    with gr.Column():
      with gr.Box():
        radio_misclassified = gr.Radio(["Yes", "No"], label="Do you want to view Misclassified images?")
        slider_misclassified_num_images = gr.Slider(minimum=1, maximum =10, value = 1, step =1, visible= False, interactive = False)
        button_misclassified = gr.Button("View Misclassified Output", visible = False)
        # txt_misclassified = gr.Textbox ("Misclassified output here" , visible = True)
        output_gallery_misclassification=gr.Gallery(label="Misclassification Output (Predicted/Truth)", min_width=512,columns=5, visible = False)

  radio_gradcam.change(fn=view_gradcam_images, inputs=radio_gradcam, outputs=[slider_gradcam_num_images, dropdown_gradcam_layer,slider_gradcam_opacity,button_gradcam, output_gallery_gradcam])
  button_gradcam.click(fn = process_gradcam_images, inputs = [slider_gradcam_num_images,dropdown_gradcam_layer,slider_gradcam_opacity], outputs = [output_gallery_gradcam,output_gallery_gradcam])

  radio_misclassified.change(fn=view_misclassified_images, inputs=radio_misclassified, outputs=[slider_misclassified_num_images,button_misclassified,output_gallery_misclassification])
  button_misclassified.click(fn = process_misclassified_images, inputs = [slider_misclassified_num_images], outputs = [output_gallery_misclassification,output_gallery_misclassification])
  button_classify.click(fn=classify_image, inputs =[input_image_classify,slider_classify_num_classes,checkbox_gradcam_classify,dropdown_gradcam_classify_layer,slider_gradcam_classify_opacity], outputs = [label_classify,gallery_gradcam_classify,gallery_gradcam_classify])
demo.launch ()