danhtran2mind commited on
Commit
584e3d6
·
verified ·
1 Parent(s): 8827262

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -6,6 +6,8 @@ import requests
6
 
7
  from models.auto_encoder_gray2color import SpatialAttention
8
 
 
 
9
  # Load the saved model once at startup
10
  load_model_path = "./ckpts/best_model.h5"
11
  if not os.path.exists(load_model_path):
@@ -28,10 +30,10 @@ loaded_autoencoder = tf.keras.models.load_model(
28
  def process_image(input_img):
29
  # Convert PIL Image to numpy array and normalize
30
  img = input_img.convert("RGB")
31
- img = img.resize((256, 256)) # adjust size as needed
32
  img_array = tf.keras.preprocessing.image.img_to_array(img) / 255.0
33
  img_array = img_array[None, ...] # add batch dimension
34
-
35
  # Run inference
36
  output_array = loaded_autoencoder.predict(img_array)
37
  output_img = tf.keras.preprocessing.image.array_to_img(output_array[0])
 
6
 
7
  from models.auto_encoder_gray2color import SpatialAttention
8
 
9
+ WIDTH, HEIGHT = 512, 512
10
+
11
  # Load the saved model once at startup
12
  load_model_path = "./ckpts/best_model.h5"
13
  if not os.path.exists(load_model_path):
 
30
  def process_image(input_img):
31
  # Convert PIL Image to numpy array and normalize
32
  img = input_img.convert("RGB")
33
+ img = img.resize((WIDTH, HEIGHT)) # adjust size as needed
34
  img_array = tf.keras.preprocessing.image.img_to_array(img) / 255.0
35
  img_array = img_array[None, ...] # add batch dimension
36
+ print("img_array shpae: ", img_array.shape)
37
  # Run inference
38
  output_array = loaded_autoencoder.predict(img_array)
39
  output_img = tf.keras.preprocessing.image.array_to_img(output_array[0])