Update app.py

app.py

@@ -42,15 +42,14 @@ def inference(input_img, num_gradcam_images=1, target_layer_number=-1, transpare
         rgb_img = np.transpose(org_img, (1, 2, 0))
         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 = (num_gradcam_images - 1) // num_rows + 1
-    fig, axes = plt.subplots(num_rows, num_cols, figsize=(12, 6))
-    for i, visualization_img in enumerate(visualization):
-        row_index = i // num_cols
-        col_index = i % num_cols
-        ax = axes[row_index, col_index] if num_rows > 1 else axes[col_index]
-        ax.imshow(visualization_img)
+    num_cols = 5
+    for i in range(len(visualizations)):
+        ax = fig.add_subplot(num_rows, num_cols, i + 1)
+        ax.imshow(visualizations[i])
         ax.axis('off')
+
     plt.tight_layout()
     buffer = BytesIO()
     plt.savefig(buffer, format='png')