Update app.py

app.py

@@ -29,22 +29,14 @@ def plot_image(image_array, scale):
     plt.gca().add_patch(Rectangle((x0, x0), scale*128, scale*128, linewidth=1, edgecolor='w', facecolor='none'))
     
     plt.axis('off')
-    fig.set_size_inches((4,4))
-    with col1:
-        st.pyplot()
+    with col1: st.pyplot()
 
 # Define function to plot the prediction
-def plot_prediction(image_array, pred):
+def plot_prediction(pred):
     plt.figure(figsize=(4, 4))
-    # plt.subplot(1, 2, 1)
-    # plt.imshow(image_array, origin="lower")
-    # plt.axis('off')
-
-    # plt.subplot(1, 2, 2)
     plt.imshow(pred, origin="lower")
     plt.axis('off')
-    with col2:
-        st.pyplot()
+    with col2: st.pyplot()
 
 def cut(data0, wcs0, scale=1):
     shape = data0.shape[0]
@@ -80,24 +72,25 @@ if uploaded_file is not None:
 
         plot_image(np.log10(data+1), scale)
 
-        if st.button('Detect cavities'):                
-            data, wcs = cut(data, wcs, scale=scale)
+        with col2:
+            if st.button('Detect cavities'):                
+                data, wcs = cut(data, wcs, scale=scale)
+        
+                image_data = np.log10(data+1)
     
-            image_data = np.log10(data+1)
-
-            y_pred = 0
-            for j in [0,1,2,3]:
-                rotated = np.rot90(image_data, j)
-                pred = model.predict(rotated.reshape(1, 128, 128, 1)).reshape(128 ,128)
-                pred = np.rot90(pred, -j)
-                y_pred += pred / 4
-
-            # 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")
+                y_pred = 0
+                for j in [0,1,2,3]:
+                    rotated = np.rot90(image_data, j)
+                    pred = model.predict(rotated.reshape(1, 128, 128, 1)).reshape(128 ,128)
+                    pred = np.rot90(pred, -j)
+                    y_pred += pred / 4
+    
+                # ccd = CCDData(pred, unit="adu", wcs=wcs)
+                # ccd.write(f"predicted.fits", overwrite=True)
+    
+                plot_prediction(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")