Update app.py

app.py

@@ -85,7 +85,12 @@ if uploaded_file is not None:
                 pred = np.rot90(pred, -j)
                 y_pred += pred / 4
 
-            # 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_prediction(image_data, y_pred)
+
+            if st.button('Download FITS File'):
+                with open('predicted.fits', 'rb') as f:
+                    data = f.read()
+                st.download_button(label="Download", data=data, file_name="predicted.fits", mime="application/octet-stream")