Update app.py

app.py

@@ -68,7 +68,8 @@ def plot_decomposed(decomposed):
     for i in range(N):
         new = np.where(decomposed == i+1, 1, 0)
         x0, y0 = center_of_mass(new)
-        plt.text(y0, x0, f"{i+1}", fontsize=15, color="white")
+        color = "white" if i < N//2 else "black"
+        plt.text(y0, x0, f"{i+1}", ha="center", va="center", fontsize=15, color=color)
     
     plt.axis('off')
     with colC: st.pyplot()
@@ -183,11 +184,11 @@ if uploaded_file is not None:
         cavs = decompose_cavity(y_pred)
 
         ccd = CCDData(y_pred, unit="adu", wcs=wcs)
-        ccd.write(f"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"predicted_{i+1}.fits", overwrite=True)
+            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")
@@ -197,15 +198,11 @@ if uploaded_file is not None:
     except: image_decomposed = np.zeros((128,128))
     plot_decomposed(image_decomposed)
 
-    # with col6:
-    #     try:
-    #         name = 'predictions.zip'
-    #         with open(name, 'rb') as f:
-    #             res = f.read()
-    #     except:
-    #         name = uploaded_file
-    #         with open(name, 'rb') as f:
-    #             res = f.read()
+    shutil.make_archive("predictions.zip", 'zip', "predictions")
+
+    with col6:
+        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=name, mime="application/octet-stream")
+        download = st.download_button(label="Download", data=res, file_name='predictions.zip', mime="application/octet-stream")