Update app.py

app.py

@@ -6,6 +6,8 @@ import requests
 
 from models.auto_encoder_gray2color import SpatialAttention
 
+WIDTH, HEIGHT = 512, 512
+
 # Load the saved model once at startup
 load_model_path = "./ckpts/best_model.h5"
 if not os.path.exists(load_model_path):
@@ -28,10 +30,10 @@ loaded_autoencoder = tf.keras.models.load_model(
 def process_image(input_img):
     # Convert PIL Image to numpy array and normalize
     img = input_img.convert("RGB")
-    img = img.resize((256, 256))  # adjust size as needed
+    img = img.resize((WIDTH, HEIGHT))  # adjust size as needed
     img_array = tf.keras.preprocessing.image.img_to_array(img) / 255.0
     img_array = img_array[None, ...]  # add batch dimension
-
+    print("img_array shpae: ", img_array.shape)
     # Run inference
     output_array = loaded_autoencoder.predict(img_array)
     output_img = tf.keras.preprocessing.image.array_to_img(output_array[0])