Update app.py

app.py

@@ -32,6 +32,7 @@ def plot_image(image, scale):
     plt.imshow(image, origin="lower")
     plt.gca().add_patch(Rectangle((x0-0.5, x0-0.5), scale*128, scale*128, linewidth=1, edgecolor='w', facecolor='none'))
     plt.axis('off')
+    plt.tight_layout()
     with colA: st.pyplot()
 
 # Define function to plot the prediction
@@ -183,63 +184,65 @@ with col:
 if uploaded_file is not None:
     data, wcs = load_file(uploaded_file)
 
-if "data" in locals():    
-    # Make six columns for buttons
-    _, col1, col2, col3, col4, col5, col6, _ = st.columns([bordersize,0.5,0.5,0.5,0.5,0.5,0.5,bordersize])
-    col1.subheader("Input image")
-    col3.subheader("Prediction")
-    col5.subheader("Decomposed")
-    col6.subheader("")
+if "data" not in locals():
+    data = np.zeros((128,128))
+
+# Make six columns for buttons
+_, col1, col2, col3, col4, col5, col6, _ = st.columns([bordersize,0.5,0.5,0.5,0.5,0.5,0.5,bordersize])
+col1.subheader("Input image")
+col3.subheader("Prediction")
+col5.subheader("Decomposed")
+col6.subheader("")
+
+with col1:
+    st.markdown("""<style>[data-baseweb="select"] {margin-top: -46px;}</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", on_change=change_scale)
+    scale = int(scale.split("x")[0]) // 128
+
+# Detect button
+with col3: detect = st.button('Detect', key="detect")
+
+# Threshold slider
+with col4:
+    st.markdown("")
+    # st.markdown("""<style>[data-baseweb="select"] {margin-top: -36px;}</style>""", unsafe_allow_html=True)
+    threshold = st.slider("Threshold", 0.0, 1.0, 0.0, 0.05, key="threshold") #, label_visibility="hidden")
     
-    with col1:
-        st.markdown("""<style>[data-baseweb="select"] {margin-top: -46px;}</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", on_change=change_scale)
-        scale = int(scale.split("x")[0]) // 128
-
-    # Detect button
-    with col3: detect = st.button('Detect', key="detect")
-
-    # Threshold slider
-    with col4:
-        st.markdown("")
-        # st.markdown("""<style>[data-baseweb="select"] {margin-top: -36px;}</style>""", unsafe_allow_html=True)
-        threshold = st.slider("Threshold", 0.0, 1.0, 0.0, 0.05, key="threshold") #, label_visibility="hidden")
-        
-    # Decompose button
-    with col5: decompose = st.button('Decompose', key="decompose")
-        
-    # Make two columns for plots
-    _, colA, colB, colC, _ = st.columns([bordersize,1,1,1,bordersize])
-    
-    image = np.log10(data+1)
-    plot_image(image, scale)
+# Decompose button
+with col5: decompose = st.button('Decompose', key="decompose")
     
-    if detect or threshold:
-    # if st.session_state.get("detect", True):
-        y_pred, wcs = cut_n_predict(data, wcs, scale)
-        
-        y_pred_th = np.where(y_pred > threshold, y_pred, 0)
-                
-        plot_prediction(y_pred_th)
+# Make two columns for plots
+_, colA, colB, colC, _ = st.columns([bordersize,1,1,1,bordersize])
+
+image = np.log10(data+1)
+plot_image(image, scale)
 
-        if decompose or st.session_state.get("download", False):            
-            image_decomposed = decompose_cavity(y_pred_th)
+if detect or threshold:
+# if st.session_state.get("detect", True):
+    y_pred, wcs = cut_n_predict(data, wcs, scale)
     
-            plot_decomposed(image_decomposed)
+    y_pred_th = np.where(y_pred > threshold, y_pred, 0)
+            
+    plot_prediction(y_pred_th)
 
-            with col6:
-                st.markdown("<br style='margin:4px 0'>", 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")
+    if decompose or st.session_state.get("download", False):            
+        image_decomposed = decompose_cavity(y_pred_th)
+
+        plot_decomposed(image_decomposed)
+
+        with col6:
+            st.markdown("<br style='margin:4px 0'>", 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")