Update app.py

app.py

@@ -3,6 +3,8 @@ from huggingface_hub import from_pretrained_keras
 model = from_pretrained_keras("Plsek/CADET-v1")
 
 # Basic libraries
+import os
+import shutil
 import numpy as np
 import matplotlib.pyplot as plt
 from matplotlib.patches import Rectangle
@@ -20,19 +22,20 @@ 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")
 
 bordersize = 0.6
 _, col, _ = st.columns([bordersize, 3, bordersize])
 
+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.<br> To use this tool: upload your image, select the scale of interest, and make a prediction!<br> If you use this tool for your research, please cite [Plšek et al. 2023](https://arxiv.org/abs/2304.05457)")
-    
+    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("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)).")
     
     # Create file uploader widget
@@ -183,17 +186,17 @@ if uploaded_file is not None:
         for i, cav in enumerate(cavs):
             ccd = CCDData(cav, unit="adu", wcs=wcs)
             ccd.write(f"predicted_{i+1}.fits", overwrite=True)
+            image_decomposed += (i+1) * np.where(cav > 0, 1, 0)
 
-            image_decomposed += (i+1) * cav
-
+        shutil.make_archive("predictions", 'zip', "predictions")
         np.save("decomposed.npy", image_decomposed)
 
     try: image_decomposed = np.load("decomposed.npy")
     except: image_decomposed = np.zeros((128,128))
     plot_decomposed(image_decomposed)
 
-    #     with col4:
-    #         pass
-    #         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")
+    with col6:
+        pass
+        # 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.zip", mime="application/octet-stream")