Update app.py

app.py

@@ -40,7 +40,7 @@ def resize_image_pil(image, new_width, new_height):
     return resized
 
 
-def inference(input_img, transparency=0.5, target_layer_number=-1, grad_cam_option="Yes"):
+def inference(input_img, transparency=0.5, target_layer_number=-1, grad_cam_option="Yes",top_classes=3):
     input_img = resize_image_pil(input_img, 32, 32)
     input_img = np.array(input_img)
     org_img = input_img
@@ -52,6 +52,7 @@ def inference(input_img, transparency=0.5, target_layer_number=-1, grad_cam_opti
     softmax = torch.nn.Softmax(dim=0)
     o = softmax(outputs.flatten())
     confidences = {classes[i]: float(o[i]) for i in range(10)}
+    confidences = dict(list(confidences.items())[:top_classes])
     _, prediction = torch.max(outputs, 1)
     target_layers = [model.layer2[target_layer_number]]
     cam = GradCAM(model=model, target_layers=target_layers)
@@ -77,12 +78,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.Slider(-2, -1, value=-2, step=1, label="Which Layer?"),
-        gr.Dropdown(["Yes", "No"], label="Want to see Grad Cam Images?")
+        gr.Dropdown(["Yes", "No"], label="Want to see Grad Cam Images?"),
+        gr.Number(value=3, minimum=1,maximum=10)
     ], 
     outputs=[
         "text", 
         gr.Image(width=256, height=256, label="Output"),
-        gr.Label(num_top_classes=3)
+        gr.Label()
     ],
     title=title,
     description=description,