Update app.py

app.py

@@ -25,6 +25,12 @@ 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))
@@ -135,23 +141,17 @@ def decompose_cavity(pred, th2=0.7, amin=6):
     
     return image_decomposed
 
-# @st.cache
-# def load_file(fname):
-#     with fits.open(fname) as hdul:
-#         data = hdul[0].data
-#         wcs = WCS(hdul[0].header)
-#     return data, wcs
+@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"]
 
-
-# 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")
+# os.system("mkdir -p predictions")
 
 with col:
     # Create heading and description
@@ -161,10 +161,12 @@ 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")
@@ -187,10 +189,10 @@ with col:
 #     data = np.zeros((128,128))
 
 if uploaded_file is not None:
-    # data, wcs = load_file(uploaded_file)
-    with fits.open(uploaded_file) as hdul:
-        data = hdul[0].data
-        wcs = WCS(hdul[0].header)
+    data, wcs = load_file(uploaded_file)
+    # with fits.open(uploaded_file) as hdul:
+    #     data = hdul[0].data
+    #     wcs = WCS(hdul[0].header)
     
     # 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])