Update app.py

app.py

@@ -44,16 +44,14 @@ def plot_image(image, scale):
     plt.gca().add_patch(Rectangle((x0, x0), scale*128, scale*128, linewidth=1, edgecolor='w', facecolor='none'))
     
     plt.axis('off')
-    # with colA: 
-    st.pyplot()
+    with colA: st.pyplot()
 
 # Define function to plot the prediction
 def plot_prediction(pred):
     plt.figure(figsize=(4, 4))
     plt.imshow(pred, origin="lower")
     plt.axis('off')
-    # with colB: 
-    st.pyplot()
+    with colB: st.pyplot()
 
 # Cut input image and rebin it to 128x128 pixels
 def cut(data0, wcs0, scale=1):
@@ -136,8 +134,7 @@ if uploaded_file is not None:
     _, colA, colB, colC, _ = st.columns([bordersize,1,1,1,bordersize])
     
     image = np.log10(data+1)
-    with colA:
-        plot_image(image, scale)
+    plot_image(image, scale)
     
     if detect:
         data, wcs = cut(data, wcs, scale=scale)