Update app.py

app.py

@@ -31,27 +31,25 @@ loaded_autoencoder = tf.keras.models.load_model(
 def process_image(input_img):
     # Convert PIL Image to grayscale and resize
     img = input_img.convert("L")  # Convert to grayscale (single channel)
-    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, ..., 0:1]  # Add batch dimension and keep single channel
-    # print("img_array shape: ", img_array.shape)
+    img = img.resize((WIDTH, HEIGHT))  # Resize to 512x512
+    img_array = tf.keras.preprocessing.image.img_to_array(img) / 255.0  # Normalize to [0, 1]
+    img_array = img_array[None, ..., 0:1]  # Add batch dimension, shape: (1, 512, 512, 1)
 
-    # Run inference
-    output_array = loaded_autoencoder.predict(img_array)
+    # Run inference (assuming loaded_autoencoder predicts a*b* channels)
+    output_array = loaded_autoencoder.predict(img_array)  # Shape: (1, 512, 512, 2) for a*b*
     print("output_array shape: ", output_array.shape)
 
-    # Assuming output_array has shape (1, 512, 512, 2) for U and V channels
-    # Extract Y (grayscale input) and UV (model output)
-    y_channel = img_array[0, :, :, 0]  # Grayscale input (Y channel)
-    uv_channels = output_array[0]  # Model output (U and V channels)
+    # Extract L* (grayscale input) and a*b* (model output)
+    L_channel = img_array[0, :, :, 0] * 100.0  # Denormalize L* to [0, 100]
+    ab_channels = output_array[0] * 128.0  # Denormalize a*b* to [-128, 128]
 
-    # Combine Y, U, V into a 3-channel YUV image
-    yuv_image = np.stack([y_channel, uv_channels[:, :, 0], uv_channels[:, :, 1]], axis=-1)
+    # Combine L*, a*, b* into a 3-channel L*a*b* image
+    lab_image = np.stack([L_channel, ab_channels[:, :, 0], ab_channels[:, :, 1]], axis=-1)  # Shape: (512, 512, 3)
 
-    # Convert YUV to RGB
-    yuv_image = yuv_image * 255.0  # Denormalize
-    rgb_image = Image.fromarray(yuv_image.astype(np.uint8), mode="YCbCr")  # Use YCbCr (alias for YUV in PIL)
-    rgb_image = rgb_image.convert("RGB")  # Convert to RGB
+    # Convert L*a*b* to RGB
+    rgb_array = lab2rgb(lab_image)  # Convert to RGB, output in [0, 1]
+    rgb_array = np.clip(rgb_array, 0, 1) * 255.0  # Scale to [0, 255]
+    rgb_image = Image.fromarray(rgb_array.astype(np.uint8), mode="RGB")  # Create RGB PIL image
 
     return rgb_image