Update app.py

app.py

@@ -34,8 +34,7 @@ def inference(input_img, num_gradcam_images=1, target_layer_number=-1, transpare
     _, prediction = torch.max(outputs, 1)
 
     visualization =[]
-    #target_layers = [[modellayers[-1]], [modellayers[-2]]]
-    for item in range(1, num_gradcam_images):
+    for item in range(1, num_gradcam_images+1):
         cam = GradCAM(model=model, target_layers = [model.layer2[-item]])
         grayscale_cam = cam(input_tensor=input_img, targets=None)
         grayscale_cam = grayscale_cam[0, :]
@@ -43,10 +42,8 @@ def inference(input_img, num_gradcam_images=1, target_layer_number=-1, transpare
         visualization.append(show_cam_on_image(org_img/255, grayscale_cam, use_rgb=True, image_weight=transparency))
 
     fig = plt.figure(figsize=(12, 5))
-    num_rows = 2
-    num_cols = 5
     for i in range(len(visualization)):
-        ax = fig.add_subplot(num_rows, num_cols, i + 1)
+        ax = fig.add_subplot(2, 5, i + 1)
         ax.imshow(visualization[i])
         ax.axis('off')
 
@@ -101,14 +98,14 @@ examples = [["cat.jpg", 1, -1, 0.8, True, 3, 3],
 demo = gr.Interface(
     inference,
     inputs=[gr.Image(width=256, height=256, label="Input Image"),
-            gr.Slider(1, 10, value=1, step=1, label="Number of GradCAM Images"),
+            gr.Slider(1, 2, value=1, step=1, label="Number of GradCAM Images"),
             gr.Slider(-2, -1, value=-1, step=1, label="Which Layer?"),
             gr.Slider(0, 1, value=0.8, label="Opacity of GradCAM"),
             gr.Checkbox(value=True, label="Show Misclassified Images"),
             gr.Slider(2, 10, value=3, step=1, label="Top Predictions"),
             gr.Slider(1, 10, value=3, step=1, label="Misclassified Images")],
     outputs=[gr.Label(label="Top Predictions"),
-             gr.Image(label="Output",width=128, height=128),
+             gr.Image(label="Output",width=640, height=360),
              gr.Image(label="Misclassified Images",width=640, height=360)],
     title=title,
     description=description,