Added Top N classes

app.py

@@ -39,7 +39,7 @@ def resize_image_pil(image, new_width, new_height):
 
     return resized
 
-def inference(input_img, transparency = 0.5, is_grad_cam=True, target_layer_number = -1):
+def inference(input_img, transparency = 0.5, is_grad_cam=True, target_layer_number = -1, top_predictions=3):
     input_img = resize_image_pil(input_img, 32, 32)
     
     input_img = np.array(input_img)
@@ -62,7 +62,14 @@ def inference(input_img, transparency = 0.5, is_grad_cam=True, target_layer_numb
         visualization = show_cam_on_image(org_img/255, grayscale_cam, use_rgb=True, image_weight=transparency)
     else:
         visualization = None
-    return classes[prediction[0].item()], visualization, confidences
+
+    # Sort the confidences dictionary based on confidence values
+    sorted_confidences = dict(sorted(confidences.items(), key=lambda item: item[1], reverse=True))
+    
+    # Pick the top n predictions
+    top_n_confidences = dict(list(sorted_confidences.items())[:top_predictions])
+    
+    return classes[prediction[0].item()], visualization, top_n_confidences
 
 title = "CIFAR10 trained on ResNet18 Model with GradCAM"
 description = "A simple Gradio interface to infer on ResNet model, and get GradCAM results"
@@ -73,12 +80,13 @@ demo = gr.Interface(
         gr.Image(width=256, height=256, label="Input Image"),
         gr.Slider(0, 1, value = 0.5, label="Overall Opacity of Image"), 
         gr.Checkbox(label="Show GradCAM"),
-        gr.Slider(-2, -1, value = -2, step=1, label="Which Layer?")
+        gr.Slider(-2, -1, value = -2, step=1, label="Which Layer?"),
+        gr.Slider(2, 10, value=3, step=1, label="Number of Top Classes")
         ], 
     outputs = [
         "text", 
         gr.Image(width=256, height=256, label="Output"),
-        gr.Label(num_top_classes=3)
+        gr.Label(label="Top Classes")
         ],
     title = title,
     description = description,