Update app.py

app.py

@@ -11,8 +11,6 @@ from matplotlib.patches import Rectangle
 from astropy.io import fits
 from astropy.wcs import WCS
 from astropy.nddata import Cutout2D, CCDData
-from astropy.convolution import Gaussian2DKernel as Gauss
-from astropy.convolution import convolve
 
 # Scikit-learn
 from sklearn.cluster import DBSCAN
@@ -23,7 +21,6 @@ st.set_option('deprecation.showPyplotGlobalUse', False)
 st.set_page_config(page_title="Cavity Detection Tool", layout="wide")
 
 # HuggingFace Hub
-# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
 from huggingface_hub import from_pretrained_keras
 # from tensorflow.keras.models import load_model
 
@@ -49,17 +46,15 @@ def plot_prediction(pred):
 def plot_decomposed(decomposed):
     plt.figure(figsize=(4, 4))
     plt.imshow(decomposed, origin="lower")
-
     N = int(np.max(decomposed))
     for i in range(N):
         new = np.where(decomposed == i+1, 1, 0)
         x0, y0 = center_of_mass(new)
         color = "white" if i < N//2 else "black"
         plt.text(y0, x0, f"{i+1}", ha="center", va="center", fontsize=15, color=color)
-    
     plt.axis('off')
     with colC: st.pyplot()
-        
+
 # Define function to cut input image and rebin it to 128x128 pixels
 def cut(data0, wcs0, scale=1):
     shape = data0.shape[0]
@@ -138,6 +133,7 @@ def decompose_cavity(pred, fname, th2=0.7, amin=10):
     
     return image_decomposed
 
+# Define function that loads FITS file and return data & wcs
 @st.cache #_data
 def load_file(fname):
     with fits.open(fname) as hdul:
@@ -145,14 +141,17 @@ def load_file(fname):
         wcs = WCS(hdul[0].header)
     return data, wcs
 
+# Define function to load model
 @st.cache(allow_output_mutation=True) #_resource
 def load_CADET():
     model = from_pretrained_keras("Plsek/CADET-v1")
     # model = load_model("CADET.hdf5")
     return model
 
+
 def reset_threshold():
-    del st.session_state["threshold"]
+    # del st.session_state["threshold"]
+    st.session_state['threshold'] = 0
 
 
 # Load model
@@ -178,7 +177,7 @@ with col:
 # _, 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']) #, on_change=reset_threshold)
+    uploaded_file = st.file_uploader("Choose a FITS file", type=['fits'], on_change=reset_threshold)
 
     # with col_2:
     #     st.markdown("### Examples")