Update app.py

app.py

@@ -56,6 +56,13 @@ def plot_prediction(pred):
     plt.axis('off')
     with colB: st.pyplot()
 
+# Define function to plot the decomposed prediction
+def plot_prediction(pred):
+    plt.figure(figsize=(4, 4))
+    plt.imshow(pred, origin="lower")
+    plt.axis('off')
+    with colC: st.pyplot()
+        
 # Cut input image and rebin it to 128x128 pixels
 def cut(data0, wcs0, scale=1):
     shape = data0.shape[0]
@@ -105,8 +112,6 @@ def decompose_cavity(pred, th2=0.7, amin=10):
         cavities.append(img)
 
     return cavities
-
-np.save("pred.npy", np.zeros((128,128)))
     
 # If file is uploaded, read in the data and plot it
 if uploaded_file is not None:
@@ -157,9 +162,9 @@ if uploaded_file is not None:
 
         np.save("pred.npy", y_pred)
 
-    y_pred = np.load("pred.npy")    
+    try: y_pred = np.load("pred.npy")
+    except: y_pred = np.zeros((128,128))
     y_pred = np.where(y_pred > threshold, y_pred, 0)
-
     np.save("thresh.npy", y_pred)
         
     plot_prediction(y_pred)
@@ -177,7 +182,9 @@ if uploaded_file is not None:
         for i, cav in enumerate(cavs):
             ccd = CCDData(cav, unit="adu", wcs=wcs)
             ccd.write(f"predicted_{i+1}.fits", overwrite=True)
-        
+
+    plot_decomposed(np.zeros((128,128)))
+
     #     with col4:
     #         pass
     #         st.markdown("""<style>[data-baseweb="select"] {margin-top: 16px;}</style>""", unsafe_allow_html=True)