Update app.py

app.py

@@ -9,8 +9,12 @@ from astropy.wcs import WCS
 from astropy.nddata import Cutout2D, CCDData
 
 # Tensorflow
-from tensorflow.keras.models import load_model
-model = load_model("CADET.hdf5")
+# from tensorflow.keras.models import load_model
+# model = load_model("CADET.hdf5")
+
+# HuggingFace 
+from huggingface_hub import from_pretrained_keras
+model = from_pretrained_keras("Plsek/CADET-v1")
 
 # Streamlit
 import streamlit as st
@@ -75,15 +79,16 @@ if uploaded_file is not None:
     # Make two columns
     col1, col2, col3, col4 = st.columns(4)
     col1.subheader("Input image")
-    col3.subheader("CADET prediction")
+    col3.subheader("Prediction")
 
     # Add a slider to change the scale
     with col1:
         smooth = st.button("Smooth")
     
     with col2:
-        st.markdown("""<style>[data-baseweb="select"] {margin-top: -52px;}</style>""", unsafe_allow_html=True)
-    
+        # st.markdown("""<style>[data-baseweb="select"] {margin-top: -52px;}</style>""", unsafe_allow_html=True)
+        st.markdown("""<style>[data-baseweb="select"] {margin-top: 52px;}</style>""", unsafe_allow_html=True)
+        
         max_scale = int(data.shape[0] // 128)
         # scale = st.slider("Scale", 1, max_scale, 1, 1)
         scale = int(st.selectbox('Scale:',[i+1 for i in range(max_scale)], label_visibility="hidden"))
@@ -92,6 +97,7 @@ if uploaded_file is not None:
         detect = st.button('Detect cavities')
     
     with col4:
+        st.markdown("""<style>[data-baseweb="select"] {margin-top: 52px;}</style>""", unsafe_allow_html=True)
         download = st.button('Download')
 
 
@@ -117,9 +123,9 @@ if uploaded_file is not None:
 
         plot_prediction(y_pred)
 
-    if download:
-        ccd = CCDData(y_pred, unit="adu", wcs=wcs)
-        ccd.write("predicted.fits", overwrite=True)
-        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")
+        if download:
+            ccd = CCDData(y_pred, unit="adu", wcs=wcs)
+            ccd.write("predicted.fits", overwrite=True)
+            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")