Update app.py

app.py

@@ -95,10 +95,6 @@ if uploaded_file is not None:
     
     with col3:
         detect = st.button('Detect cavities')
-    
-    with col4:
-        st.markdown("""<style>[data-baseweb="select"] {margin-top: 32px;}</style>""", unsafe_allow_html=True)
-        download = st.button('Download')
 
     colA, colB = st.columns(2)
 
@@ -120,10 +116,14 @@ if uploaded_file is not None:
         # Thresholding
         y_pred = np.where(y_pred > 0.4, y_pred, 0)
 
-        plot_prediction(y_pred)
+        # plot_prediction(y_pred)
 
         ccd = CCDData(y_pred, unit="adu", wcs=wcs)
         ccd.write("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")
+        
+        with col4:
+            st.markdown("""<style>[data-baseweb="select"] {margin-top: 32px;}</style>""", unsafe_allow_html=True)
+            # download = st.button('Download')
+            download = st.download_button(label="Download", data=data, file_name="predicted.fits", mime="application/octet-stream")