Update app.py

app.py

@@ -8,7 +8,7 @@ from tensorflow.keras.models import load_model
 
 st.set_option('deprecation.showPyplotGlobalUse', False)
 
-st.title("FITS Image Viewer")
+st.title("Cavity Detection Tool")
 
 model = load_model("CADET.hdf5")
 
@@ -25,23 +25,44 @@ def plot_image(image_array, pred):
     st.pyplot()
 
 def cut(data0, wcs0, scale=1):
-   shape = data0.shape[0]
-   x0 = shape / 2
-   size = 128 * scale
-   cutout = Cutout2D(data0, (x0, x0), (size, size), wcs=wcs0)
-   return cutout.data, cutout.wcs
+    shape = data0.shape[0]
+    x0 = shape / 2
+    size = 128 * scale
+    cutout = Cutout2D(data0, (x0, x0), (size, size), wcs=wcs0)
+    data, wcs = cutout.data, cutout.wcs
+
+    # REGRID DATA
+    factor = size // 128
+    data = data.reshape(128, factor, 128, factor).mean(-1).mean(1)
+    
+    # REGIRD WCS
+    ra, dec = wcs.wcs_pix2world(np.array([[63, 63]]),0)[0]
+    wcs.wcs.cdelt[0] = wcs.wcs.cdelt[0] * factor
+    wcs.wcs.cdelt[1] = wcs.wcs.cdelt[1] * factor
+    wcs.wcs.crval[0] = ra
+    wcs.wcs.crval[1] = dec
+    wcs.wcs.crpix[0] = 64 / factor
+    wcs.wcs.crpix[1] = 64 / factor
+
+   return data, wcs
 
 # 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)
-        data, wcs = cut(data, wcs, scale=1)
+        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)
+
         plot_image(image_data, pred)