Update app.py

app.py

@@ -3,7 +3,7 @@ import numpy as np
 import matplotlib.pyplot as plt
 from astropy.io import fits
 from astropy.wcs import WCS
-from astropy.nddata import Cutout2D
+from astropy.nddata import Cutout2D, CCDData
 from tensorflow.keras.models import load_model
 
 st.set_option('deprecation.showPyplotGlobalUse', False)
@@ -13,7 +13,14 @@ st.title("Cavity Detection Tool")
 model = load_model("CADET.hdf5")
 
 # Define function to plot the uploaded image
-def plot_image(image_array, pred):
+def plot_image(image_array, scale):
+    plt.figure(figsize=(5, 5))
+    # plt.subplot(1, 2, 1)
+    plt.imshow(image_array, origin="lower")
+    plt.axis('off')
+
+# Define function to plot the prediction
+def plot_prediction(image_array, pred):
     plt.figure(figsize=(10, 5))
     plt.subplot(1, 2, 1)
     plt.imshow(image_array, origin="lower")
@@ -57,12 +64,17 @@ if uploaded_file is not None:
     with fits.open(uploaded_file) as hdul:
         data = hdul[0].data
         wcs = WCS(hdul[0].header)
-        data, wcs = cut(data, wcs, scale=scale)
 
-        image_data = np.log10(data+1)
-        pred = model.predict(image_data.reshape(1, 128, 128, 1)).reshape(128 ,128)
+        plot_image(np.log10(data+1), scale)
+
+        if st.button('Detect Cavity'):                
+            data, wcs = cut(data, wcs, scale=scale)
+    
+            image_data = np.log10(data+1)
+
+            pred = model.predict(image_data.reshape(1, 128, 128, 1)).reshape(128 ,128)
 
-        ccd = CCDData(pred, unit="adu", wcs=wcs)
-        ccd.write(f"predicted.fits", overwrite=True)
+            # ccd = CCDData(pred, unit="adu", wcs=wcs)
+            # ccd.write(f"predicted.fits", overwrite=True)
 
-        plot_image(image_data, pred)
+            plot_prediction(image_data, pred)