Update app.py

app.py

@@ -79,41 +79,52 @@ if uploaded_file is not None:
     col2.subheader("CADET prediction")
     
     with col1:
-        st.markdown(
-            """<style>[data-baseweb="select"] {margin-top: -50px;}</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"))
-
+        col1_1, col1_2 = st.columns(2)
+
+        with col1_2:
+            st.Button("Smooth")
+
+        with col1_1:
+            st.markdown(
+                """<style>[data-baseweb="select"] {margin-top: -50px;}</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"))
+            
     # Add a slider to change the scale
     with col1:
         plot_image(np.log10(data+1), scale)
 
     with col2:
-        if st.button('Detect cavities'):
-            data, wcs = cut(data, wcs, scale=scale)
+        col2_1, col2_2 = st.columns(2)
 
-            image = np.log10(data+1)
-            
-            y_pred = 0
-            for j in [0,1,2,3]:
-                rotated = np.rot90(image, j)
-                pred = model.predict(rotated.reshape(1, 128, 128, 1)).reshape(128 ,128)
-                pred = np.rot90(pred, -j)
-                y_pred += pred / 4
+        with col2_1:
+            if st.button('Detect cavities'):
+                data, wcs = cut(data, wcs, scale=scale)
     
-            # Thresholding
-            y_pred = np.where(y_pred > 0.4, y_pred, 0)
-
-            plot_prediction(y_pred)
-    
-            # ccd = CCDData(pred, unit="adu", wcs=wcs)
-            # ccd.write(f"predicted.fits", overwrite=True)
-            
-            # 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")
+                image = np.log10(data+1)
+                
+                y_pred = 0
+                for j in [0,1,2,3]:
+                    rotated = np.rot90(image, j)
+                    pred = model.predict(rotated.reshape(1, 128, 128, 1)).reshape(128 ,128)
+                    pred = np.rot90(pred, -j)
+                    y_pred += pred / 4
+        
+                # Thresholding
+                y_pred = np.where(y_pred > 0.4, y_pred, 0)
+
+                with col2:
+                    plot_prediction(y_pred)
+        
+                ccd = CCDData(pred, unit="adu", wcs=wcs)
+                ccd.write(f"predicted.fits", overwrite=True)
+
+                with col2_2:
+                    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")