Update app.py

app.py

@@ -33,15 +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')
-    st.pyplot()
+    with col1: 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 col2: st.pyplot()
 
 # Cut input image and rebin it to 128x128 pixels
 def cut(data0, wcs0, scale=1):
@@ -77,54 +76,46 @@ if uploaded_file is not None:
     col1, col2 = st.columns(2)
     col1.subheader("Input image")
     col2.subheader("CADET prediction")
-    
-    with col1:
-        col1_1, col1_2 = st.columns(2)
-
-        with col1_2:
-            st.Button("Smooth")
 
-        with col1_1:
-            st.markdown(
-                """<style>[data-baseweb="select"] {margin-top: -50px;}</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"))
-            
     # Add a slider to change the scale
     with col1:
-        plot_image(np.log10(data+1), scale)
+        st.markdown(
+            """<style>[data-baseweb="select"] {margin-top: -50px;}</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:
-        col2_1, col2_2 = st.columns(2)
+        detect = st.button('Detect cavities'):
 
-        with col2_1:
-            if st.button('Detect cavities'):
-                data, wcs = cut(data, wcs, scale=scale)
-    
-                image = np.log10(data+1)
-                
-                y_pred = 0
-                for j in [0,1,2,3]:
-                    rotated = np.rot90(image, j)
-                    pred = model.predict(rotated.reshape(1, 128, 128, 1)).reshape(128 ,128)
-                    pred = np.rot90(pred, -j)
-                    y_pred += pred / 4
-        
-                # Thresholding
-                y_pred = np.where(y_pred > 0.4, y_pred, 0)
+    if detect:
+        data, wcs = cut(data, wcs, scale=scale)
 
-                with col2:
-                    plot_prediction(y_pred)
+        image = np.log10(data+1)
         
-                ccd = CCDData(pred, unit="adu", wcs=wcs)
-                ccd.write(f"predicted.fits", overwrite=True)
-
-                with col2_2:
-                    if st.button('Download FITS File'):
-                        with open('predicted.fits', 'rb') as f:
-                            data = f.read()
-                        st.download_button(label="Download", data=data, file_name="predicted.fits", mime="application/octet-stream")
+        y_pred = 0
+        for j in [0,1,2,3]:
+            rotated = np.rot90(image, j)
+            pred = model.predict(rotated.reshape(1, 128, 128, 1)).reshape(128 ,128)
+            pred = np.rot90(pred, -j)
+            y_pred += pred / 4
+
+        # Thresholding
+        y_pred = np.where(y_pred > 0.4, y_pred, 0)
+
+        plot_prediction(y_pred)
+
+    # with col2:
+    #     download = st.button('Download FITS File')
+
+    # if download:
+    #     ccd = CCDData(pred, unit="adu", wcs=wcs)
+    #     ccd.write(f"predicted.fits", overwrite=True)
+    #     with open('predicted.fits', 'rb') as f:
+    #         data = f.read()
+    #     st.download_button(label="Download", data=data, file_name="predicted.fits", mime="application/octet-stream")