Update app.py

app.py

@@ -27,36 +27,33 @@ from huggingface_hub import from_pretrained_keras
 
 # Define function to plot the uploaded image
 def plot_image(image, scale):
-    # plt.figure(figsize=(4, 4))
-    # x0 = image.shape[0] // 2 - scale * 128 / 2
-    # plt.imshow(image, origin="lower")
-    # plt.gca().add_patch(Rectangle((x0-0.5, x0-0.5), scale*128, scale*128, linewidth=1, edgecolor='w', facecolor='none'))
-    # plt.axis('off')
-    # plt.tight_layout()
-    # with colA: st.pyplot()
-    with colA: st.image(image)
+    plt.figure(figsize=(4, 4))
+    x0 = image.shape[0] // 2 - scale * 128 / 2
+    plt.imshow(image, origin="lower")
+    plt.gca().add_patch(Rectangle((x0-0.5, x0-0.5), scale*128, scale*128, linewidth=1, edgecolor='w', facecolor='none'))
+    plt.axis('off')
+    plt.tight_layout()
+    with colA: st.pyplot()
 
 # Define function to plot the prediction
 def plot_prediction(pred):
-    # plt.figure(figsize=(4, 4))
-    # plt.imshow(pred, origin="lower", norm=Normalize(vmin=0, vmax=1))
-    # plt.axis('off')
-    # with colB: st.pyplot()
-    with colB: st.image(pred)
+    plt.figure(figsize=(4, 4))
+    plt.imshow(pred, origin="lower", norm=Normalize(vmin=0, vmax=1))
+    plt.axis('off')
+    with colB: st.pyplot()
 
 # Define function to plot the decomposed prediction
 def plot_decomposed(decomposed):
-    # plt.figure(figsize=(4, 4))
-    # plt.imshow(decomposed, origin="lower")
-    # N = int(np.max(decomposed))
-    # for i in range(N):
-    #     new = np.where(decomposed == i+1, 1, 0)
-    #     x0, y0 = center_of_mass(new)
-    #     color = "white" if i < N//2 else "black"
-    #     plt.text(y0, x0, f"{i+1}", ha="center", va="center", fontsize=15, color=color)
-    # plt.axis('off')
-    # with colC: st.pyplot()
-    with colC: st.image(decomposed)
+    plt.figure(figsize=(4, 4))
+    plt.imshow(decomposed, origin="lower")
+    N = int(np.max(decomposed))
+    for i in range(N):
+        new = np.where(decomposed == i+1, 1, 0)
+        x0, y0 = center_of_mass(new)
+        color = "white" if i < N//2 else "black"
+        plt.text(y0, x0, f"{i+1}", ha="center", va="center", fontsize=15, color=color)
+    plt.axis('off')
+    with colC: st.pyplot()
 
 # Define function to cut input image and rebin it to 128x128 pixels
 def cut(data0, wcs0, scale=1):
@@ -170,12 +167,13 @@ os.system("rm -R -- */")
 # os.system("mkdir -p predictions")
 
 with col:
-    # Create heading and description
-    st.markdown("<h1 align='center'>Cavity Detection Tool</h1>", unsafe_allow_html=True)    
-    st.markdown("Cavity Detection Tool (CADET) is a machine learning pipeline trained to detect X-ray cavities from noisy Chandra images of early-type galaxies.")
-    st.markdown("To use this tool: upload your image, select the scale of interest, make a prediction, and decompose it into individual cavities!")
-    st.markdown("Input images should be in units of counts, centred at the galaxy center, and point sources should be filled with surrounding background ([dmfilth](https://cxc.cfa.harvard.edu/ciao/ahelp/dmfilth.html)).")
-    st.markdown("If you use this tool for your research, please cite [Plšek et al. 2023](https://arxiv.org/abs/2304.05457)")
+    with st.container:
+        # Create heading and description
+        st.markdown("<h1 align='center'>Cavity Detection Tool</h1>", unsafe_allow_html=True)    
+        st.markdown("Cavity Detection Tool (CADET) is a machine learning pipeline trained to detect X-ray cavities from noisy Chandra images of early-type galaxies.")
+        st.markdown("To use this tool: upload your image, select the scale of interest, make a prediction, and decompose it into individual cavities!")
+        st.markdown("Input images should be in units of counts, centred at the galaxy center, and point sources should be filled with surrounding background ([dmfilth](https://cxc.cfa.harvard.edu/ciao/ahelp/dmfilth.html)).")
+        st.markdown("If you use this tool for your research, please cite [Plšek et al. 2023](https://arxiv.org/abs/2304.05457)")
 
 # _, col_1, col_2, col_3, _ = st.columns([bordersize, 2.0, 0.5, 0.5, bordersize])