Update app.py

app.py

@@ -7,6 +7,8 @@ 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
 
 # HuggingFace
 from huggingface_hub import from_pretrained_keras
@@ -25,16 +27,23 @@ st.markdown("Cavity Detection Tool (CADET) is a machine learning pipeline traine
 uploaded_file = st.file_uploader("Choose a FITS file", type=['fits'])
 
 # Define function to plot the uploaded image
-def plot_image(image_array, scale):
+def plot_image(image, scale):
     plt.figure(figsize=(4, 4))
 
-    x0 = image_array.shape[0] // 2 - scale * 128 / 2
-    plt.imshow(image_array, origin="lower")
+    x0 = image.shape[0] // 2 - scale * 128 / 2
+    plt.imshow(image, origin="lower")
     plt.gca().add_patch(Rectangle((x0, x0), scale*128, scale*128, linewidth=1, edgecolor='w', facecolor='none'))
     
     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))
@@ -83,7 +92,7 @@ if uploaded_file is not None:
     
     with col2:
         # st.markdown("""<style>[data-baseweb="select"] {margin-top: -52px;}</style>""", unsafe_allow_html=True)
-        st.markdown("""<style>[data-baseweb="select"] {margin-top: 32px;}</style>""", unsafe_allow_html=True)
+        st.markdown("""<style>[data-baseweb="select"] {margin-top: 22px;}</style>""", unsafe_allow_html=True)
         
         max_scale = int(data.shape[0] // 128)
         # scale = st.slider("Scale", 1, max_scale, 1, 1)
@@ -94,12 +103,12 @@ if uploaded_file is not None:
 
     colA, colB = st.columns(2)
 
-    plot_image(np.log10(data+1), scale)
-
+    image = np.log10(data+1)
+    if smooth: image = smooth_image(image, scale)
+    plot_image(image, scale)
     
     if detect:
         data, wcs = cut(data, wcs, scale=scale)
-
         image = np.log10(data+1)
         
         y_pred = 0
@@ -112,14 +121,14 @@ if uploaded_file is not None:
         # Thresholding
         y_pred = np.where(y_pred > 0.4, y_pred, 0)
 
-        # plot_prediction(y_pred)
+        plot_prediction(y_pred)
 
         ccd = CCDData(y_pred, unit="adu", wcs=wcs)
         ccd.write("predicted.fits", overwrite=True)
         with open('predicted.fits', 'rb') as f:
             data = f.read()
         
-        # with col4:
-        #     st.markdown("""<style>[data-baseweb="select"] {margin-top: 32px;}</style>""", unsafe_allow_html=True)
+        with col4:
+            st.markdown("""<style>[data-baseweb="select"] {margin-top: 32px;}</style>""", unsafe_allow_html=True)
         #     # download = st.button('Download')
         #     download = st.download_button(label="Download", data=data, file_name="predicted.fits", mime="application/octet-stream")