Update app.py

app.py

@@ -15,9 +15,10 @@ model = load_model("CADET.hdf5")
 
 # Define function to plot the uploaded image
 def plot_image(image_array, scale):
-    plt.figure(figsize=(5, 5))
+    st.set_plot_config(plt, figsize=(4, 4))
+    plt.figure(figsize=(4, 4))
     # plt.subplot(1, 2, 1)
-    x0 = image_array.shape[0] // 2
+    x0 = image_array.shape[0] // 2 - scale * 128 / 64
     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')
@@ -25,7 +26,8 @@ def plot_image(image_array, scale):
 
 # Define function to plot the prediction
 def plot_prediction(image_array, pred):
-    plt.figure(figsize=(10, 5))
+    st.set_plot_config(plt, figsize=(8, 4))
+    plt.figure(figsize=(8, 4))
     plt.subplot(1, 2, 1)
     plt.imshow(image_array, origin="lower")
     plt.axis('off')