Update app.py

app.py

@@ -93,24 +93,27 @@ if uploaded_file is not None:
         plot_image(np.log10(data+1), scale)
 
     with col2:
-        st.button('Detect cavities')      
-        data, wcs = cut(data, wcs, scale=scale)
-
-        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
-
-        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")
+        if st.button('Detect cavities', disable=True):
+            data, wcs = cut(data, wcs, scale=scale)
+    
+            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)
+            
+            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")