Update app.py

app.py

@@ -27,24 +27,6 @@ st.set_option('deprecation.showPyplotGlobalUse', False)
 st.set_page_config(page_title="Cavity Detection Tool", layout="wide")
 # st.title("Cavity Detection Tool")
 
-bordersize = 0.6
-_, col, _ = st.columns([bordersize, 3, bordersize])
-
-os.system("rm -r predictions")
-os.system("rm predictions.zip Views")
-os.system("mkdir predictions")
-
-with col:
-    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)).")
-    st.markdown("If you use this tool for your research, please cite [Plšek et al. 2023](https://arxiv.org/abs/2304.05457)")
-    
-    # Create file uploader widget
-    uploaded_file = st.file_uploader("Choose a FITS file", type=['fits'])
-
 # Define function to plot the uploaded image
 def plot_image(image, scale):
     plt.figure(figsize=(4, 4))
@@ -125,6 +107,27 @@ def decompose_cavity(pred, th2=0.7, amin=10):
         cavities.append(img)
 
     return cavities
+
+
+
+
+bordersize = 0.6
+_, col, _ = st.columns([bordersize, 3, bordersize])
+
+# os.system("rm -r predictions")
+# os.system("rm predictions.zip Views")
+# os.system("mkdir -p predictions")
+
+with col:
+    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)).")
+    st.markdown("If you use this tool for your research, please cite [Plšek et al. 2023](https://arxiv.org/abs/2304.05457)")
+    
+    # Create file uploader widget
+    uploaded_file = st.file_uploader("Choose a FITS file", type=['fits'])
     
 # If file is uploaded, read in the data and plot it
 if uploaded_file is not None:
@@ -171,40 +174,44 @@ if uploaded_file is not None:
             pred = np.rot90(pred, -j)
             y_pred += pred / 4
 
-        np.save("pred.npy", y_pred)
+        # np.save("pred.npy", y_pred)
 
-    try: y_pred = np.load("pred.npy")
+    # try: y_pred = np.load("pred.npy")
+    # except: y_pred = np.zeros((128,128))
+    try: y_pred
     except: y_pred = np.zeros((128,128))
     y_pred = np.where(y_pred > threshold, y_pred, 0)
-    np.save("thresh.npy", y_pred)
+    # np.save("thresh.npy", y_pred)
         
     plot_prediction(y_pred)
 
     if decompose:
-        y_pred = np.load("thresh.npy")
+        # y_pred = np.load("thresh.npy")
         
         cavs = decompose_cavity(y_pred)
 
         ccd = CCDData(y_pred, unit="adu", wcs=wcs)
-        ccd.write(f"predictions/predicted.fits", overwrite=True)
+        # 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)
+            # 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)
 
         # shutil.make_archive("predictions.zip", 'zip', "predictions")
-        np.save("decomposed.npy", image_decomposed)
+        # np.save("decomposed.npy", image_decomposed)
 
-    try: image_decomposed = np.load("decomposed.npy")
+    # try: image_decomposed = np.load("decomposed.npy")
+    # except: image_decomposed = np.zeros((128,128))
+    try: image_decomposed
     except: image_decomposed = np.zeros((128,128))
     plot_decomposed(image_decomposed)
 
-    shutil.make_archive("predictions", 'zip', "predictions")
+    # shutil.make_archive("predictions", 'zip', "predictions")
 
     with col6:
-        with open('predictions.zip', 'rb') as f:
-            res = f.read()
+        # with open('predictions.zip', 'rb') as f:
+        #     res = f.read()
     #     # 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='predictions.zip', mime="application/octet-stream")
+        download = st.download_button(label="Download", data=ccd, file_name='prediction.fits', mime="application/octet-stream")