Update app.py

app.py

@@ -44,14 +44,16 @@ def plot_image(image, scale):
     plt.gca().add_patch(Rectangle((x0, x0), scale*128, scale*128, linewidth=1, edgecolor='w', facecolor='none'))
     
     plt.axis('off')
-    with colA: st.pyplot()
+    # with colA: 
+    st.pyplot()
 
 # Define function to plot the prediction
 def plot_prediction(pred):
     plt.figure(figsize=(4, 4))
     plt.imshow(pred, origin="lower")
     plt.axis('off')
-    with colB: st.pyplot()
+    # with colB: 
+    st.pyplot()
 
 # Cut input image and rebin it to 128x128 pixels
 def cut(data0, wcs0, scale=1):
@@ -133,26 +135,26 @@ if uploaded_file is not None:
     _, colA, colB, colC, _ = st.columns([bordersize,1,1,1,bordersize])
     
     image = np.log10(data+1)
-    plot_image(image, scale)
-
-    # with colB:
-    #     threshold = st.slider("", 0.0, 1.0, 0.4, 0.05, label_visibility="hidden")
+    with colA:
+        plot_image(image, scale)
     
-    # if detect:
-    #     data, wcs = cut(data, wcs, scale=scale)
-    #     image = np.log10(data+1)
+    if detect:
+        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
-
-    #     # Thresholding
-    #     y_pred = np.where(y_pred > threshold, y_pred, 0)
-                
-    #     plot_prediction(y_pred)
+        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
+
+        # Thresholding
+        y_pred = np.where(y_pred > threshold, y_pred, 0)
+
+        with colB:
+            threshold = st.slider("", 0.0, 1.0, 0.4, 0.05, label_visibility="hidden")
+            plot_prediction(y_pred)
 
         # if decompose:
         #     cavs = decompose_cavity(y_pred, )