Update app.py

app.py

@@ -180,15 +180,15 @@ if uploaded_file is not None:
     # except: y_pred = np.zeros((128,128))
     try: y_pred
     except: y_pred = np.zeros((128,128))
-    y_pred = np.where(y_pred > threshold, y_pred, 0)
+    y_pred_th = np.where(y_pred > threshold, y_pred, 0)
     # np.save("thresh.npy", y_pred)
         
-    plot_prediction(y_pred)
+    plot_prediction(y_pred_th)
 
     if decompose:
         # y_pred = np.load("thresh.npy")
         
-        cavs = decompose_cavity(y_pred)
+        cavs = decompose_cavity(y_pred_th)
 
         # ccd = CCDData(y_pred, unit="adu", wcs=wcs)
         # ccd.write(f"predictions/predicted.fits", overwrite=True)
@@ -209,9 +209,9 @@ if uploaded_file is not None:
 
     # 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:
+    #     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")