Update app.py

app.py

@@ -4,7 +4,7 @@ import shutil
 import numpy as np
 from scipy.ndimage import center_of_mass
 import matplotlib.pyplot as plt
-from matplotlib.colors import Normalize
+from matplotlib.colors import Normalize, LogNorm
 from matplotlib.patches import Rectangle
 
 # Astropy
@@ -29,7 +29,7 @@ from huggingface_hub import from_pretrained_keras
 def plot_image(image, scale):
     plt.figure(figsize=(4, 4))
     x0 = image.shape[0] // 2 - scale * 128 / 2
-    plt.imshow(image, origin="lower")
+    plt.imshow(image, origin="lower", norm=LogNorm())
     plt.gca().add_patch(Rectangle((x0-0.5, x0-0.5), scale*128, scale*128, linewidth=1, edgecolor='w', facecolor='none'))
     plt.axis('off')
     plt.tight_layout()
@@ -219,17 +219,14 @@ if uploaded_file is not None:
     
     MIN = np.min(np.where(data == 0, 1, data))
     if MIN < 1: data = data / MIN
-    image = np.log10(data+1)
 
 if st.session_state.get("example", True):
     fname = "NGC4649_example"
     os.system(f'mkdir -p {fname}')
     data, wcs = load_file(f"{fname}.fits")
-    image = np.log10(data+1)
 
 if "data" not in locals():
     data = np.zeros((128,128))
-    image = data
         
 # Make six 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])
@@ -262,7 +259,7 @@ _, colA, colB, colC, _ = st.columns([bordersize,1,1,1,bordersize])
 
 if "data" in locals():
     # NORMALIZE IMAGE
-    plot_image(image, scale)
+    plot_image(data, scale)
     
     if detect or threshold or st.session_state.get("decompose", False):