Update app.py

app.py

@@ -207,28 +207,24 @@ if uploaded_file is not None:
             
             cavs = decompose_cavity(y_pred_th)
     
-            # ccd = CCDData(y_pred, unit="adu", wcs=wcs)
-            # ccd.write(f"predictions/predicted.fits", overwrite=True)
+            ccd = CCDData(y_pred, unit="adu", wcs=wcs)
+            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)
-    
-        # try: image_decomposed = np.load("decomposed.npy")
-        # except: image_decomposed = np.zeros((128,128))
+            shutil.make_archive("predictions.zip", 'zip', "predictions")
+
+            shutil.make_archive("predictions", 'zip', "predictions")
+        
         try: _ = image_decomposed
         except: image_decomposed = np.zeros((128,128))
         plot_decomposed(image_decomposed)
     
-        # shutil.make_archive("predictions", 'zip', "predictions")
-    
-        # with col6:
-        #     ccd = CCDData(y_pred, unit="adu", wcs=wcs)
-        #     # 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.download_button(label="Download", data=ccd, file_name='prediction.fits', mime="application/octet-stream")
+        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.download_button(label="Download", data=res, file_name='prediction.fits', mime="application/octet-stream")