Update app.py

app.py

@@ -11,31 +11,39 @@ st.set_option('deprecation.showPyplotGlobalUse', False)
 
 st.title("Cavity Detection Tool")
 
+st.text("Cavity Detection Tool (CADET) is a machine learning pipeline trained to detect X-ray cavities from noisy Chandra images of early-type galaxies.\
+To use this tool upload your image, select the scale of interest and make a prediction! If you use output of this tool in your research please cite [Plšek et al. 2023](https://arxiv.org/abs/2304.05457)")
+
 model = load_model("CADET.hdf5")
 
+# Create file uploader widget
+uploaded_file = st.file_uploader("Choose a FITS file", type=['fits'])
+
+col1, col2 = st.columns(2)
+
 # Define function to plot the uploaded image
 def plot_image(image_array, scale):
-    # st.set_plot_config(plt, figsize=(4, 4))
     plt.figure(figsize=(4, 4))
-    # plt.subplot(1, 2, 1)
+
     x0 = image_array.shape[0] // 2 - scale * 128 / 2
     plt.imshow(image_array, origin="lower")
     plt.gca().add_patch(Rectangle((x0, x0), scale*128, scale*128, linewidth=1, edgecolor='w', facecolor='none'))
+    
     plt.axis('off')
-    st.pyplot(width=200)
+    fig.set_size_inches((4,4))
+    st.pyplot()
 
 # Define function to plot the prediction
 def plot_prediction(image_array, pred):
-    # st.set_plot_config(plt, figsize=(8, 4))
-    plt.figure(figsize=(8, 4))
-    plt.subplot(1, 2, 1)
-    plt.imshow(image_array, origin="lower")
-    plt.axis('off')
+    plt.figure(figsize=(4, 4))
+    # plt.subplot(1, 2, 1)
+    # plt.imshow(image_array, origin="lower")
+    # plt.axis('off')
 
-    plt.subplot(1, 2, 2)
+    # plt.subplot(1, 2, 2)
     plt.imshow(pred, origin="lower")
     plt.axis('off')
-    st.pyplot(width=400)
+    col2.pyplot()
 
 def cut(data0, wcs0, scale=1):
     shape = data0.shape[0]
@@ -59,8 +67,6 @@ def cut(data0, wcs0, scale=1):
 
     return data, wcs
 
-# Create file uploader widget
-uploaded_file = st.file_uploader("Choose a FITS file", type=['fits'])
 
 # If file is uploaded, read in the data and plot it
 if uploaded_file is not None: