Update app.py

app.py

@@ -81,6 +81,22 @@ def cut(data0, wcs0, scale=1):
 
     return data, wcs
 
+
+@st.cache_data
+def cut_n_predict(data, scale):
+    data, wcs = cut(data, wcs, scale=scale)
+    image = np.log10(data+1)
+    
+    y_pred = 0
+    for j in [0,1,2,3]:
+        rotated = np.rot90(image, j)
+        pred = model.predict(rotated.reshape(1, 128, 128, 1)).reshape(128 ,128)
+        pred = np.rot90(pred, -j)
+        y_pred += pred / 4
+
+    return y_pred
+
+# Define function to decomposed prediction into cavities
 def decompose_cavity(pred, th2=0.7, amin=10):
     X, Y = pred.nonzero()
     data = np.array([X,Y]).reshape(2, -1)
@@ -159,61 +175,53 @@ if uploaded_file is not None:
     
     image = np.log10(data+1)
     plot_image(image, scale)
-    
-    if detect:
 
-        with col4:
+
+    with col4:
             st.markdown("""<style>[data-baseweb="select"] {margin-top: -36px;}</style>""", unsafe_allow_html=True)
             threshold = st.slider("", 0.0, 1.0, 0.0, 0.05, label_visibility="hidden")
+    
+    if detect:
+        y_pred = predict(data, scale)
         
-        data, wcs = cut(data, wcs, scale=scale)
-        image = np.log10(data+1)
-        
-        y_pred = 0
-        for j in [0,1,2,3]:
-            rotated = np.rot90(image, j)
-            pred = model.predict(rotated.reshape(1, 128, 128, 1)).reshape(128 ,128)
-            pred = np.rot90(pred, -j)
-            y_pred += pred / 4
-
         # np.save("pred.npy", y_pred)
     
-        # try: y_pred = np.load("pred.npy")
-        # except: y_pred = np.zeros((128,128))
-        try: y_pred
-        except: y_pred = np.zeros((128,128))
-        y_pred_th = np.where(y_pred > threshold, y_pred, 0)
-        # np.save("thresh.npy", y_pred)
+    # try: y_pred = np.load("pred.npy")
+    # except: y_pred = np.zeros((128,128))
+    try: y_pred
+    except: y_pred = np.zeros((128,128))
+    y_pred_th = np.where(y_pred > threshold, y_pred, 0)
+    # np.save("thresh.npy", y_pred)
             
-        plot_prediction(y_pred_th)
-    
-        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)
+    plot_prediction(y_pred_th)
     
-            # shutil.make_archive("predictions.zip", 'zip', "predictions")
-            # np.save("decomposed.npy", image_decomposed)
-    
-        # try: image_decomposed = np.load("decomposed.npy")
-        # except: image_decomposed = np.zeros((128,128))
-        try: image_decomposed
-        except: image_decomposed = np.zeros((128,128))
-        plot_decomposed(image_decomposed)
-    
-        # 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")
+    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)
+
+        # shutil.make_archive("predictions.zip", 'zip', "predictions")
+        # np.save("decomposed.npy", image_decomposed)
+
+    # try: image_decomposed = np.load("decomposed.npy")
+    # except: image_decomposed = np.zeros((128,128))
+    try: image_decomposed
+    except: image_decomposed = np.zeros((128,128))
+    plot_decomposed(image_decomposed)
+
+    # 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")