Update app.py

app.py

@@ -18,11 +18,11 @@ def plot_image(image_array, scale):
     # st.set_plot_config(plt, figsize=(4, 4))
     plt.figure(figsize=(4, 4))
     # plt.subplot(1, 2, 1)
-    x0 = image_array.shape[0] // 2 - scale * 128 / 64
+    x0 = image_array.shape[0] // 2 - scale * 128 / 2
     plt.imshow(image_array, origin="lower")
     plt.gca().add_patch(Rectangle((x0, x0), scale*128, scale*128, linewidth=1, edgecolor='w', facecolor='none'))
     plt.axis('off')
-    st.pyplot(width=400)
+    st.pyplot(width=200)
 
 # Define function to plot the prediction
 def plot_prediction(image_array, pred):
@@ -35,7 +35,7 @@ def plot_prediction(image_array, pred):
     plt.subplot(1, 2, 2)
     plt.imshow(pred, origin="lower")
     plt.axis('off')
-    st.pyplot(width=800)
+    st.pyplot(width=400)
 
 def cut(data0, wcs0, scale=1):
     shape = data0.shape[0]
@@ -73,14 +73,19 @@ if uploaded_file is not None:
 
         plot_image(np.log10(data+1), scale)
 
-        if st.button('Detect Cavity'):                
+        if st.button('Detect cavities'):                
             data, wcs = cut(data, wcs, scale=scale)
     
             image_data = np.log10(data+1)
 
-            pred = model.predict(image_data.reshape(1, 128, 128, 1)).reshape(128 ,128)
+            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
 
             # ccd = CCDData(pred, unit="adu", wcs=wcs)
             # ccd.write(f"predicted.fits", overwrite=True)
 
-            plot_prediction(image_data, pred)
+            plot_prediction(image_data, y_pred)