Update app.py

app.py

@@ -33,14 +33,14 @@ def plot_image(image_array, scale):
     plt.gca().add_patch(Rectangle((x0, x0), scale*128, scale*128, linewidth=1, edgecolor='w', facecolor='none'))
     
     plt.axis('off')
-    with col1: 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 col2: st.pyplot()
+    with colB: st.pyplot()
 
 # Cut input image and rebin it to 128x128 pixels
 def cut(data0, wcs0, scale=1):
@@ -73,26 +73,33 @@ if uploaded_file is not None:
         wcs = WCS(hdul[0].header)
     
     # Make two columns
-    col1, col2 = st.columns(2)
+    col1, col2, col3, col4 = st.columns(4)
     col1.subheader("Input image")
-    col2.subheader("CADET prediction")
+    col3.subheader("CADET prediction")
 
     # Add a slider to change the scale
     with col1:
-        st.markdown(
-            """<style>[data-baseweb="select"] {margin-top: -52px;}</style>""",
-            unsafe_allow_html=True
-        )
+        smooth = st.button("Smooth")
+    
+    with col2:
+        st.markdown("""<style>[data-baseweb="select"] {margin-top: -52px;}</style>""", unsafe_allow_html=True)
     
         max_scale = int(data.shape[0] // 128)
         # scale = st.slider("Scale", 1, max_scale, 1, 1)
         scale = int(st.selectbox('Scale:',[i+1 for i in range(max_scale)], label_visibility="hidden"))
-            
-    plot_image(np.log10(data+1), scale)
-
-    with col2:
+    
+    with col3:
         detect = st.button('Detect cavities')
+    
+    with col4:
+        download = st.button('Download')
 
+
+    colA, colB = st.columns(2)
+
+    plot_image(np.log10(data+1), scale)
+
+    
     if detect:
         data, wcs = cut(data, wcs, scale=scale)
 
@@ -110,11 +117,8 @@ if uploaded_file is not None:
 
         plot_prediction(y_pred)
 
-    with col2:
-        download = st.button('Download')
-
     if download:
-        ccd = CCDData(pred, unit="adu", wcs=wcs)
+        ccd = CCDData(y_pred, unit="adu", wcs=wcs)
         ccd.write("predicted.fits", overwrite=True)
         with open('predicted.fits', 'rb') as f:
             data = f.read()