Update app.py

app.py

@@ -123,7 +123,15 @@ def decompose_cavity(pred, th2=0.7, amin=10):
 
         cavities.append(img)
 
-    return cavities
+    ccd = CCDData(pred, unit="adu", wcs=wcs)
+    ccd.write(f"predictions/predicted.fits", overwrite=True)
+    image_decomposed = np.zeros((128,128))
+    for i, cav in enumerate(cavities):
+        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)
+
+    return image_decomposed
 
 # @st.cache
 # def zip_predictions():
@@ -204,23 +212,23 @@ if uploaded_file is not None:
         if decompose:
             # y_pred = np.load("thresh.npy")
             
-            cavs = decompose_cavity(y_pred_th)
-    
-            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)
-                image_decomposed += (i+1) * np.where(cav > 0, 1, 0)
-            
-        try: _ = image_decomposed
-        except: image_decomposed = np.zeros((128,128))
-        plot_decomposed(image_decomposed)
+            image_decomposed = decompose_cavity(y_pred_th)
     
-    with col6:
-        shutil.make_archive("predictions", 'zip', "predictions")
-        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")
+            # 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)
+            #     image_decomposed += (i+1) * np.where(cav > 0, 1, 0)
+
+            try: _ = image_decomposed
+            except: image_decomposed = np.zeros((128,128))
+            plot_decomposed(image_decomposed)
+        
+            with col6:
+                shutil.make_archive("predictions", 'zip', "predictions")
+                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.zip', mime="application/octet-stream")