Update app.py

app.py

@@ -94,7 +94,7 @@ def cut_n_predict(data, scale):
         pred = np.rot90(pred, -j)
         y_pred += pred / 4
 
-    return y_pred
+    return y_pred, wcs
 
 # Define function to decomposed prediction into cavities
 def decompose_cavity(pred, th2=0.7, amin=10):
@@ -182,7 +182,7 @@ if uploaded_file is not None:
             threshold = st.slider("", 0.0, 1.0, 0.0, 0.05, label_visibility="hidden")
     
     if detect:
-        y_pred = cut_n_predict(data, scale)
+        y_pred, wcs = cut_n_predict(data, wcs, scale)
         
         # np.save("pred.npy", y_pred)