Update app.py

app.py

@@ -33,11 +33,26 @@ def process_image(input_img):
     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)
+    # print("img_array shape: ", img_array.shape)
+
     # Run inference
     output_array = loaded_autoencoder.predict(img_array)
-    output_img = tf.keras.preprocessing.image.array_to_img(output_array[0])
-    return output_img
+    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)
+
+    # Combine Y, U, V into a 3-channel YUV image
+    yuv_image = np.stack([y_channel, uv_channels[:, :, 0], uv_channels[:, :, 1]], axis=-1)
+
+    # 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
+
+    return rgb_image
 
 custom_css = """
 body {background: linear-gradient(135deg, #232526 0%, #414345 100%) !important;}