Update app.py

app.py

@@ -166,8 +166,8 @@ if uploaded_file is not None:
     col6.subheader("")
     
     with col1:
-        # st.markdown("""<style>[data-baseweb="select"] {margin-top: -26px;}</style>""", unsafe_allow_html=True)
-        st.markdown("<style>{margin-top: -26px;}</style>") #, unsafe_allow_html=True)
+        st.markdown("""<style>[data-baseweb="select"] {margin-top: -26px;}</style>""", unsafe_allow_html=True)
+        st.markdown("<style>{margin-top: -26px;}</style>", unsafe_allow_html=True)
         max_scale = int(data.shape[0] // 128)
         scale = st.selectbox('Scale:',[f"{(i+1)*128}x{(i+1)*128}" for i in range(max_scale)], label_visibility="hidden")
         scale = int(scale.split("x")[0]) // 128
@@ -183,23 +183,6 @@ if uploaded_file is not None:
         
     # Decompose button
     with col5: decompose = st.button('Decompose', key="decompose")
-
-    with col6:
-        st.markdown("<br><br>")
-        # st.markdown("""<style>[data-baseweb="select"] {margin-top: 16px;}</style>""", unsafe_allow_html=True)
-        fname = uploaded_file.name.strip(".fits")
-
-        if st.session_state.get("decompose", False):
-            st.session_state.disabled = False
-
-        shutil.make_archive("predictions", 'zip', "predictions")
-        with open('predictions.zip', 'rb') as f:
-            res = f.read()
-        
-        download = st.download_button(label="Download", data=res, 
-                                      file_name=f'{fname}_{int(scale*128)}.zip', 
-                                      disabled=st.session_state.get("disabled", True), 
-                                      mime="application/octet-stream")
         
     # Make two columns for plots
     _, colA, colB, colC, _ = st.columns([bordersize,1,1,1,bordersize])
@@ -214,8 +197,23 @@ if uploaded_file is not None:
                 
         plot_prediction(y_pred_th)
 
-        print(type(download))
-        if decompose or download:            
+        if decompose or st.session_state.get("download", False):            
             image_decomposed = decompose_cavity(y_pred_th)
     
-            plot_decomposed(image_decomposed)
+            plot_decomposed(image_decomposed)
+
+            with col6:
+                st.markdown("<br><br>", unsafe_allow_html=True)
+                # st.markdown("""<style>[data-baseweb="select"] {margin-top: 16px;}</style>""", unsafe_allow_html=True)
+                fname = uploaded_file.name.strip(".fits")
+        
+                # if st.session_state.get("download", False):
+        
+                shutil.make_archive("predictions", 'zip', "predictions")
+                with open('predictions.zip', 'rb') as f:
+                    res = f.read()
+                
+                download = st.download_button(label="Download", data=res, key="download", 
+                                              file_name=f'{fname}_{int(scale*128)}.zip', 
+                                              # disabled=st.session_state.get("disabled", True), 
+                                              mime="application/octet-stream")