Update app.py

app.py

@@ -24,8 +24,6 @@ from sklearn.cluster import DBSCAN
 # Streamlit
 import streamlit as st
 st.set_option('deprecation.showPyplotGlobalUse', False)
-st.set_page_config(page_title="Cavity Detection Tool", layout="wide")
-# st.title("Cavity Detection Tool")
 
 # Define function to plot the uploaded image
 def plot_image(image, scale):
@@ -58,7 +56,7 @@ def plot_decomposed(decomposed):
     plt.axis('off')
     with colC: st.pyplot()
         
-# Cut input image and rebin it to 128x128 pixels
+# Define function to cut input image and rebin it to 128x128 pixels
 def cut(data0, wcs0, scale=1):
     shape = data0.shape[0]
     x0 = shape / 2
@@ -81,7 +79,7 @@ def cut(data0, wcs0, scale=1):
 
     return data, wcs
 
-
+# Define function to apply cutting and produce a prediction
 @st.cache
 def cut_n_predict(data, wcs, scale):
     data, wcs = cut(data, wcs, scale=scale)
@@ -96,13 +94,13 @@ def cut_n_predict(data, wcs, scale):
 
     return y_pred, wcs
 
-# Define function to decomposed prediction into cavities
+# Define function to decompose prediction into individual cavities
 @st.cache
 def decompose_cavity(pred, th2=0.7, amin=10):
     X, Y = pred.nonzero()
     data = np.array([X,Y]).reshape(2, -1)
 
-    # DBSCAN CLUSTERING ALGORITHM
+    # DBSCAN clustering
     try: clusters = DBSCAN(eps=1.0, min_samples=3).fit(data.T).labels_
     except: clusters = []
 
@@ -115,14 +113,16 @@ def decompose_cavity(pred, th2=0.7, amin=10):
         xi, yi = X[b], Y[b]
         img[xi, yi] = pred[xi, yi]
 
-        # THRESHOLDING #2
+        # Thresholding #2
         if not (img > th2).any(): continue
 
-        # MINIMAL AREA
+        # Minimal area
         if np.sum(img) <= amin: continue
 
         cavities.append(img)
 
+    # Save raw and decomposed predictions to predictions folder
+    os.system("mkdir -p predictions")
     ccd = CCDData(pred, unit="adu", wcs=wcs)
     ccd.write(f"predictions/predicted.fits", overwrite=True)
     image_decomposed = np.zeros((128,128))
@@ -133,24 +133,14 @@ def decompose_cavity(pred, th2=0.7, amin=10):
 
     return image_decomposed
 
-# @st.cache
-# def zip_predictions():
-#     shutil.make_archive("predictions.zip", 'zip', "predictions")
-#     with open('predictions.zip', 'rb') as f:
-#         res = f.read()
-#         return res
-
+# Use wide layout and create columns
+st.set_page_config(page_title="Cavity Detection Tool", layout="wide")
 bordersize = 0.6
 _, col, _ = st.columns([bordersize, 3, bordersize])
 
-# if os.path.exists("pred.npy"): os.system("rm pred.npy")
-# os.system("rm -r predictions")
-# os.system("rm predictions.zip Views")
-os.system("mkdir -p predictions")
-
 with col:
-    st.markdown("# Cavity Detection Tool")
-    
+    # Create heading and description
+    st.markdown("# Cavity Detection Tool")    
     st.markdown("Cavity Detection Tool (CADET) is a machine learning pipeline trained to detect X-ray cavities from noisy Chandra images of early-type galaxies.")
     st.markdown("To use this tool: upload your image, select the scale of interest, and make a prediction!")
     st.markdown("Input images should be centred at the centre of the galaxy and point sources should be filled with surrounding background ([dmfilth](https://cxc.cfa.harvard.edu/ciao/ahelp/dmfilth.html)).")
@@ -170,61 +160,46 @@ if uploaded_file is not None:
     col1.subheader("Input image")
     col3.subheader("Prediction")
     col5.subheader("Decomposed")
-    col6.subheader("")
+    # col6.subheader("")
     
     with col1:
-        
         # st.markdown("""<style>[data-baseweb="select"] {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
-    
-    with col3:
-        detect = st.button('Detect')
 
-    with col5:
-        decompose = st.button('Decompose')
+    # Detect button
+    with col3: detect = st.button('Detect')
+
+    # Threshold slider
+    with col4:
+        # 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) #, label_visibility="hidden")
+        
+    # Decompose button
+    with col5: decompose = st.button('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)
-
-
-    with col4:
-            st.markdown("""<style>[data-baseweb="select"] {margin-top: -36px;}</style>""", unsafe_allow_html=True)
-            threshold = st.slider("", 0.0, 1.0, 0.0, 0.05, label_visibility="hidden")
     
     if detect or threshold:
         y_pred, wcs = cut_n_predict(data, wcs, scale)
         
-        # np.save("pred.npy", y_pred)
-    
-        # try: y_pred = np.load("pred.npy")
+        # try: _ = y_pred
         # except: y_pred = np.zeros((128,128))
-        try: _ = y_pred
-        except: y_pred = np.zeros((128,128))
         y_pred_th = np.where(y_pred > threshold, y_pred, 0)
-        # np.save("thresh.npy", y_pred)
                 
         plot_prediction(y_pred_th)
         
         if decompose:
-            # y_pred = np.load("thresh.npy")
             
             image_decomposed = decompose_cavity(y_pred_th)
     
-            # ccd = CCDData(y_pred, unit="adu", wcs=wcs)
-            # ccd.write(f"predictions/predicted.fits", overwrite=True)
-            # image_decomposed = np.zeros((128,128))
-            # for i, cav in enumerate(cavs):
-            #     ccd = CCDData(cav, unit="adu", wcs=wcs)
-            #     ccd.write(f"predictions/predicted_{i+1}.fits", overwrite=True)
-            #     image_decomposed += (i+1) * np.where(cav > 0, 1, 0)
-
-            try: _ = image_decomposed
-            except: image_decomposed = np.zeros((128,128))
+            # try: _ = image_decomposed
+            # except: image_decomposed = np.zeros((128,128))
             plot_decomposed(image_decomposed)
         
             with col6:
@@ -235,4 +210,4 @@ if uploaded_file is not None:
                 st.markdown("")
                 # st.markdown("""<style>[data-baseweb="select"] {margin-top: 16px;}</style>""", unsafe_allow_html=True)
                 fname = uploaded_file.name.strip(".fits")    
-                download = st.download_button(label="Download", data=res, file_name=f'pred_{int(scale*128)}.zip', mime="application/octet-stream")
+                download = st.download_button(label="Download", data=res, file_name=f'{fname}_{int(scale*128)}.zip', mime="application/octet-stream")