Update app.py

app.py

@@ -105,6 +105,8 @@ def decompose_cavity(pred, th2=0.7, amin=10):
         cavities.append(img)
 
     return cavities
+
+np.save("pred.npy", y_pred)
     
 # If file is uploaded, read in the data and plot it
 if uploaded_file is not None:
@@ -157,19 +159,24 @@ if uploaded_file is not None:
 
     y_pred = np.load("pred.npy")    
     y_pred = np.where(y_pred > threshold, y_pred, 0)
+
+    np.save("thresh.npy", y_pred)
         
     plot_prediction(y_pred)
 
     # with colC:
     #     st.markdown("""<style>[data-baseweb="select"] {margin-top: -36px;}</style>""", unsafe_allow_html=True)
 
-    # if decompose:
-    #     cavs = decompose_cavity(y_pred)
+    if decompose:
+        y_pred = np.save("thresh.npy")
         
-    #     ccd = CCDData(y_pred, unit="adu", wcs=wcs)
-    #     ccd.write("predicted.fits", overwrite=True)
-    #     with open('predicted.fits', 'rb') as f:
-    #         res = f.read()
+        cavs = decompose_cavity(y_pred)
+
+        ccd = CCDData(y_pred, unit="adu", wcs=wcs)
+        ccd.write(f"predicted.fits", overwrite=True)
+        for i, cav in enumerate(cavs):
+            ccd = CCDData(cav, unit="adu", wcs=wcs)
+            ccd.write(f"predicted_{i+1}.fits", overwrite=True)
         
     #     with col4:
     #         pass