Update app.py

app.py

@@ -18,12 +18,20 @@ model = from_pretrained_keras("Plsek/CADET-v1")
 import streamlit as st
 st.set_option('deprecation.showPyplotGlobalUse', False)
 
+st.set_page_config(
+    page_title="Cavity Detection Tool",
+    # page_icon="👋",
+    layout="wide"
+    # initial_sidebar_state="expanded",
+    }
+)
 
-st.title("Cavity Detection Tool")
+
+# st.title("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. To use this tool: upload your image, select the scale of interest, and make a prediction! If you use this tool for your research, please cite [Plšek et al. 2023](https://arxiv.org/abs/2304.05457)")
 
-st.markdown("Input images should be centered 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).")
+st.markdown("Input images should be centered 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)).")
 
 
 # Create file uploader widget
@@ -40,13 +48,6 @@ def plot_image(image, scale):
     plt.axis('off')
     with colA: st.pyplot()
 
-# Define function to smooth image
-def smooth_image(image, scale):
-    smoothed = convolve(image, boundary = "wrap", nan_treatment="interpolate",
-                    kernel = Gauss(x_stddev = 2, y_stddev = 2))
-
-    return smoothed
-
 # Define function to plot the prediction
 def plot_prediction(pred):
     plt.figure(figsize=(4, 4))
@@ -77,6 +78,32 @@ def cut(data0, wcs0, scale=1):
 
     return data, wcs
 
+def decompose_cavity(pred, th2=0.7, amin=10):
+    X, Y = pred.nonzero()
+    data = np.array([X,Y]).reshape(2, -1)
+
+    # DBSCAN CLUSTERING ALGORITHM
+    try: clusters = DBSCAN(eps=1.5, min_samples=3).fit(data.T).labels_
+    except: clusters = []
+
+    N = len(set(clusters))
+    cavities = []
+
+    for i in range(N):
+        img = np.zeros((128,128))
+        b = clusters == i
+        xi, yi = X[b], Y[b]
+        img[xi, yi] = pred[xi, yi]
+
+        # THRESHOLDING #2
+        if not (img > th2).any(): continue
+
+        # MINIMAL AREA
+        if np.sum(img) <= amin: continue
+
+        cavities.append(img)
+
+    return cavities
     
 # If file is uploaded, read in the data and plot it
 if uploaded_file is not None:
@@ -88,24 +115,24 @@ if uploaded_file is not None:
     col1, col2, col3, col4 = st.columns(4)
     col1.subheader("Input image")
     col3.subheader("Prediction")
-
-    # Add a slider to change the scale
-    with col1:
-        smooth = st.button("Smooth")
     
-    with col2:
-        st.markdown("""<style>[data-baseweb="select"] {margin-top: 17px;}</style>""", unsafe_allow_html=True)
+    with col1:
+        st.markdown("""<style>[data-baseweb="select"] {margin-top: 16px;}</style>""", unsafe_allow_html=True)
         max_scale = int(data.shape[0] // 128)
-        scale = int(st.selectbox('Scale:',[i+1 for i in range(max_scale)], label_visibility="hidden"))
+        # scale = int(st.selectbox('Scale:',[i+1 for i in range(max_scale)], label_visibility="hidden"))
+        scale = int(st.selectbox('Scale:',[f"{(i+1)*128}x{(i+1)*128}" for i in range(max_scale)], label_visibility="hidden"))
+        scale = scale.split("x")[0] // 128
     
-    with col3:
+    with col2:
         detect = st.button('Detect cavities')
 
+    with col3:
+        decompose = st.button('Docompose cavities')
+        
     # Make two columns for plots
     colA, colB = st.columns(2)
 
     image = np.log10(data+1)
-    if smooth: image = smooth_image(image, scale)
     plot_image(image, scale)
     
     if detect:
@@ -122,7 +149,10 @@ if uploaded_file is not None:
         # Thresholding
         y_pred = np.where(y_pred > 0.4, y_pred, 0)
 
-        plot_prediction(y_pred)
+        # if decompose:
+        #     cavs = decompose_cavity(y_pred, )
+        
+        plot_prediction(y_pred, decompose)
 
         ccd = CCDData(y_pred, unit="adu", wcs=wcs)
         ccd.write("predicted.fits", overwrite=True)
@@ -131,6 +161,6 @@ if uploaded_file is not None:
         
         with col4:
             pass
-            # st.markdown("""<style>[data-baseweb="select"] {margin-top: 32px;}</style>""", unsafe_allow_html=True)
+            st.markdown("""<style>[data-baseweb="select"] {margin-top: 16px;}</style>""", unsafe_allow_html=True)
         #     # download = st.button('Download')
             download = st.download_button(label="Download", data=res, file_name="predicted.fits", mime="application/octet-stream")