Update app.py

app.py

@@ -8,6 +8,7 @@ from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam.utils.image import show_cam_on_image
 from resnet_lightning import ResNet18Model
 import gradio as gr
+from visualize import plot_gradcam_images, plot_misclassified_images
 
 model = ResNet18Model.load_from_checkpoint("epoch=19-step=3920.ckpt")
 
@@ -19,6 +20,19 @@ inv_normalize = transforms.Normalize(
 classes = ('plane', 'car', 'bird', 'cat', 'deer',
            'dog', 'frog', 'horse', 'ship', 'truck')
 
+
+def get_layer(layer_name):
+    if layer_name == "0":
+        return [model.prep[-1]]
+    elif layer_name == "1":
+        return [model.layer1_x[-1]]
+    elif layer_name == "2":
+        return [model.layer2_x[-1]]
+    elif layer_name == "3":
+        return [model.layer3_x[-1]]
+    else:
+        return None
+
 def resize_image_pil(image, new_width, new_height):
     img =  Image.fromarray(np.array(image))
     width, height = img.size
@@ -32,10 +46,11 @@ def resize_image_pil(image, new_width, new_height):
     return resized
 
 
-def inference(input_img, transparancy = 0.5, target_layer_number = -1):
+def inference(input_img, show_gradcam, layer_name, num_classes, transparancy = 0.5):
     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)
@@ -44,14 +59,27 @@ def inference(input_img, transparancy = 0.5, target_layer_number = -1):
     print(outputs)
     softmax =  torch.nn.Softmax(dim=0)
     o = softmax(outputs.flatten())
-    confidences = {classes[i]:float(o[i]) for i in range(10)}
+    
+    output_numpy = np.squeeze(np.asarray(outputs.numpy()))
+    index_sort = np.argsort(output_numpy)[::-1]
+    
+    confidences = {}
+    for i in range(int(num_classes)):
+        confidences[classes[index_sort[i]]] = float(o[index_sort[i]])
+
+    
     prediction= torch.max(outputs, 1)
-    target_layers = [model.layer2[target_layer_number]]
-    cam = GradCAM(model=model, target_layers=target_layers)
-    grayscale_cam = cam(input_tensor= input_img)
-    grayscale_cam = grayscale_cam[0, :]
-    visualization = show_cam_on_image(org_img/255,grayscale_cam,use_rgb=True,
+
+    if show_gradcam:
+        target_layers = get_layer(layer_name)
+        cam = GradCAM(model=model, target_layers=target_layers)
+        grayscale_cam = cam(input_tensor= input_img)
+        grayscale_cam = grayscale_cam[0, :]
+        visualization = show_cam_on_image(org_img/255,grayscale_cam,use_rgb=True,
                                      image_weight=transparancy)
+    else:
+        visualization = org_img
+
     
     return classes[int(prediction[0].item())], visualization, confidences
     
@@ -59,17 +87,19 @@ demo = gr.Interface(
     inference,
     inputs = [
         gr.Image(width=256,height=256,label="input image"),
+        gr.Number(value=3, maximum=10, minimum=1,step=1.0, precision=0,label="Number of classes to display"),
+        gr.Checkbox(True,label="Show GradCAM Image"),
+        gr.Dropdown(layer_name, value="layer3_x", label="Which layer for Gradcam"),
         gr.Slider(0,1,value=0.5,label="Overall opacity of the overelay"),
-        gr.Slider(-2,-1, value =-2, step=1, label= "Which layer for Gradcam")
     ],
     outputs = [
-        "text",
         gr.Image(width= 256, height=256,label="Output"),
-        gr.Label(num_top_classes=3)        
+        gr.Label(label="Confidences", container=True, show_label= True)        
     ],
     title = "CIFAR 10 trained on ResNet model in pytorch lightning with Gradcam",
     description = " A simple gradio inference to infer on resnet18 model",
-    examples = [["cat.jpg", 0.5, -1],["dog.jpg",0.7,-2]]
+    examples = [["cat.jpg", True, "layer3_x", 10, -1],
+                ["dog.jpg", False, "layer3_x", 4, -1]]
 )
 
 if __name__ == "__main__":