Update app.py

app.py

@@ -7,6 +7,7 @@ import os
 import shutil
 import numpy as np
 import matplotlib.pyplot as plt
+from matplotlib.colors import LogNorm
 from matplotlib.patches import Rectangle
 
 # Astropy
@@ -44,11 +45,9 @@ with col:
 # Define function to plot the uploaded image
 def plot_image(image, scale):
     plt.figure(figsize=(4, 4))
-
     x0 = image.shape[0] // 2 - scale * 128 / 2
     plt.imshow(image, origin="lower")
     plt.gca().add_patch(Rectangle((x0, x0), scale*128, scale*128, linewidth=1, edgecolor='w', facecolor='none'))
-    
     plt.axis('off')
     with colA: st.pyplot()
 
@@ -62,7 +61,14 @@ def plot_prediction(pred):
 # Define function to plot the decomposed prediction
 def plot_decomposed(decomposed):
     plt.figure(figsize=(4, 4))
-    plt.imshow(decomposed, origin="lower")
+    plt.imshow(decomposed, origin="lower", norm=LogNorm())
+
+    N = np.max(decomposed)
+    for i in range(N):
+        new = np.where(decomposed == i+1, 1, 0)
+        x0, y0 = center_of_mass(new)
+        plt.text(y0, x0, f"{i+1}", fontsize=15, color="white")
+    
     plt.axis('off')
     with colC: st.pyplot()
         
@@ -121,7 +127,6 @@ if uploaded_file is not None:
     with fits.open(uploaded_file) as hdul:
         data = hdul[0].data
         wcs = WCS(hdul[0].header)
-        y_pred = np.zeros((128,128))
     
     # Make four columns for buttons
     _, col1, col2, col3, col4, col5, col6, _ = st.columns([bordersize,0.5,0.5,0.5,0.5,0.5,0.5,bordersize])
@@ -134,7 +139,6 @@ if uploaded_file is not None:
         max_scale = int(data.shape[0] // 128)
         scale = st.selectbox('Scale:',[f"{(i+1)*128}x{(i+1)*128}" for i in range(max_scale)], label_visibility="hidden")
         scale = int(scale.split("x")[0]) // 128
-        # np.save("pred.npy", y_pred)
     
     with col3:
         detect = st.button('Detect')