Update app.py

app.py

@@ -24,9 +24,6 @@ st.markdown("Cavity Detection Tool (CADET) is a machine learning pipeline traine
 # Create file uploader widget
 uploaded_file = st.file_uploader("Choose a FITS file", type=['fits'])
 
-# Make two columns
-col1, col2 = st.columns(2)
-
 # Define function to plot the uploaded image
 def plot_image(image_array, scale):
     plt.figure(figsize=(4, 4))
@@ -69,45 +66,53 @@ def cut(data0, wcs0, scale=1):
 
     return data, wcs
 
+# Make two columns
+col1, col2 = st.columns(2)
+col1.subheader("Input image")
+col2.subheader("CADET prediction")
+
+with col1:
+    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"))
+
+with col2:
+    button == st.button('Detect cavities', disabled=True)
+
+    
 # If file is uploaded, read in the data and plot it
 if uploaded_file is not None:
-    col1.subheader("Input image")
-    col2.subheader("CADET prediction")
-    
     with fits.open(uploaded_file) as hdul:
         data = hdul[0].data
         wcs = WCS(hdul[0].header)
 
     # Add a slider to change the scale
     with col1:
-        max_scale = int(data.shape[0] // 128)
-        # scale = st.slider("Scale", 1, max_scale, 1, 1)
-
-        st.markdown(
-            """<style>[data-baseweb="select"] {margin-top: -50px;}</style>""",
-            unsafe_allow_html=True
-        )
-                
-        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:
-        if st.button('Detect cavities', disabled=True):
-            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)
-            
+    button.disabled = False
+        
+    if button:
+        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)
+
+        with col2:
             plot_prediction(y_pred)
     
             # ccd = CCDData(pred, unit="adu", wcs=wcs)