putting everything on the same plot

app.py

@@ -117,7 +117,7 @@ def predict_and_analyze(model_name, num_channels, dim, image):
     activation_2 = normalize_array(activation_2)
     
     # convert input image to RGB
-    input_image = Image.fromarray(np.uint8(cm.magma(input_image)*255))
+    input_pil_image = Image.fromarray(np.uint8(cm.magma(input_image)*255))
     
     print("Plotting")
     
@@ -126,28 +126,31 @@ def predict_and_analyze(model_name, num_channels, dim, image):
     plt.rcParams['xtick.labelsize'] = ticks
     plt.rcParams['ytick.labelsize'] = ticks
     
-    fig, axs = plt.subplots(nrows=1, ncols=2, figsize=(19, 8))
+    fig, axs = plt.subplots(nrows=3, ncols=1, figsize=(8, 28))
     
-    ax1, ax2 = axs[0], axs[1]
+    ax0, ax1, ax2 = axs[0], axs[1], axs[2]
     
+    im0 = ax1.imshow(input_image, cmap=cmap,
+                     origin=origin)    
     im1 = ax1.imshow(activation_1, cmap=cmap,
                      origin=origin)
     im2 = ax2.imshow(activation_2, cmap=cmap, 
                      origin=origin) 
     
-    ims = [im1, im2]
+    ims = [im0, im1, im2]
     
     for (i, ax) in enumerate(axs):
         divider = make_axes_locatable(ax)
         cax = divider.append_axes('right', size='5%', pad=0.05)
         fig.colorbar(ims[i], cax=cax, orientation='vertical')
         
+    ax0.set_title('Input', fontsize=titles)
     ax1.set_title('First Activation', fontsize=titles)
     ax2.set_title('Second Activation', fontsize=titles)
     
     print("Sending to Hugging Face")
     
-    return output, input_image, fig
+    return output, input_pil_image, fig
 
 
 if __name__ == "__main__":
@@ -168,7 +171,7 @@ if __name__ == "__main__":
                             show_label=True), 
                 gr.File(label="Input Data", show_label=True)],
         outputs=[gr.Textbox(lines=1, label="Prediction", show_label=True), 
-                gr.Image(label="Input Image", show_label=True), 
+                # gr.Image(label="Input Image", show_label=True), 
                 #  gr.Image(label="Activation 1", show_label=True), 
                 #  gr.Image(label="Actication 2", show_label=True)],
                 gr.Plot(label="Activations", show_label=True)