Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 import numpy as np
 import matplotlib.pyplot as plt
+from matplotlib.patches import Rectangle
 from astropy.io import fits
 from astropy.wcs import WCS
 from astropy.nddata import Cutout2D, CCDData
@@ -16,8 +17,11 @@ model = load_model("CADET.hdf5")
 def plot_image(image_array, scale):
     plt.figure(figsize=(5, 5))
     # plt.subplot(1, 2, 1)
+    x0 = image_array.shape[0] // 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()
 
 # Define function to plot the prediction
 def plot_prediction(image_array, pred):
@@ -56,15 +60,15 @@ def cut(data0, wcs0, scale=1):
 # Create file uploader widget
 uploaded_file = st.file_uploader("Choose a FITS file", type=['fits'])
 
-# Add a slider to change the scale
-scale = st.slider("Scale", 1, 4, 1, 1)
-
 # If file is uploaded, read in the data and plot it
 if uploaded_file is not None:
     with fits.open(uploaded_file) as hdul:
         data = hdul[0].data
         wcs = WCS(hdul[0].header)
 
+        # Add a slider to change the scale
+        scale = st.slider("Scale", 1, 4, 1, 1)
+
         plot_image(np.log10(data+1), scale)
 
         if st.button('Detect Cavity'):