Update app.py

app.py

@@ -25,12 +25,6 @@ from sklearn.cluster import DBSCAN
 import streamlit as st
 st.set_option('deprecation.showPyplotGlobalUse', False)
 
-# Use wide layout and create columns
-st.set_page_config(page_title="Cavity Detection Tool", layout="wide")
-bordersize = 0.45
-_, col, _ = st.columns([bordersize, 3, bordersize])
-
-
 # Define function to plot the uploaded image
 def plot_image(image, scale):
     plt.figure(figsize=(4, 4))
@@ -87,7 +81,7 @@ def cut(data0, wcs0, scale=1):
     return data, wcs
 
 # Define function to apply cutting and produce a prediction
-# @st.cache
+@st.cache
 def cut_n_predict(data, wcs, scale):
     data, wcs = cut(data, wcs, scale=scale)
     image = np.log10(data+1)
@@ -102,7 +96,7 @@ def cut_n_predict(data, wcs, scale):
     return y_pred, wcs
 
 # Define function to decompose prediction into individual cavities
-# @st.cache
+@st.cache
 def decompose_cavity(pred, th2=0.7, amin=6):
     X, Y = pred.nonzero()
     data = np.array([X,Y]).reshape(2, -1)
@@ -141,17 +135,24 @@ def decompose_cavity(pred, th2=0.7, amin=6):
     
     return image_decomposed
 
-# @st.cache
+@st.cache
 def load_file(fname):
     with fits.open(fname) as hdul:
         data = hdul[0].data
         wcs = WCS(hdul[0].header)
     return data, wcs
 
-# def change_scale():
-#     del st.session_state["threshold"]
+def change_scale():
+    del st.session_state["threshold"]
+
 
-# os.system("mkdir -p predictions")
+
+# Use wide layout and create columns
+st.set_page_config(page_title="Cavity Detection Tool", layout="wide")
+bordersize = 0.45
+_, col, _ = st.columns([bordersize, 3, bordersize])
+
+os.system("mkdir -p predictions")
 
 with col:
     # Create heading and description
@@ -161,12 +162,10 @@ with col:
     st.markdown("Input images should be in units of counts, centred at the galaxy center, and point sources should be filled with surrounding background ([dmfilth](https://cxc.cfa.harvard.edu/ciao/ahelp/dmfilth.html)).")
     st.markdown("If you use this tool for your research, please cite [Plšek et al. 2023](https://arxiv.org/abs/2304.05457)")
 
-    uploaded_file = st.file_uploader("Choose a FITS file", type=['fits'])
-    
 # _, col_1, col_2, col_3, _ = st.columns([bordersize, 2.0, 0.5, 0.5, bordersize])
 
 # with col:
-    # uploaded_file = st.file_uploader("Choose a FITS file", type=['fits'])
+    uploaded_file = st.file_uploader("Choose a FITS file", type=['fits'])
 
 # with col_2:
 #     st.markdown("### Examples")
@@ -182,71 +181,68 @@ with col:
 #     uploaded_file = "NGC5813_example.fits"
 
 # If file is uploaded, read in the data and plot it
-# if uploaded_file is not None:
-#     data, wcs = load_file(uploaded_file)
-
-# if "data" not in locals():
-#     data = np.zeros((128,128))
-
 if uploaded_file is not None:
     data, wcs = load_file(uploaded_file)
+
+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")
-        
-    # 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])
+# 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")
     
-    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 detect or threshold:
+# if st.session_state.get("detect", True):
+    y_pred, wcs = cut_n_predict(data, wcs, scale)
     
-        # if decompose or st.session_state.get("download", False):            
-        #     image_decomposed = decompose_cavity(y_pred_th)
+    y_pred_th = np.where(y_pred > threshold, y_pred, 0)
+            
+    plot_prediction(y_pred_th)
+
+    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")
     
-        #     plot_decomposed(image_decomposed)
+            # if st.session_state.get("download", False):
     
-        #     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")
+            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")