Create app.py

app.py

@@ -0,0 +1,67 @@
+import gradio as gr
+import tensorflow as tf
+import numpy as np
+from PIL import Image
+from tensorflow_examples.models.pix2pix import pix2pix
+
+OUTPUT_CHANNELS = 3
+
+generator_g = pix2pix.unet_generator(OUTPUT_CHANNELS, norm_type='instancenorm')
+generator_g.load_weights("vibrantGAN-generator-g-final.weights.h5")
+
+def preprocess_single_image(image, target_height=256, target_width=256):
+    # Convert PIL image to tensorflow tensor
+    image = tf.convert_to_tensor(np.array(image))
+
+    # Ensure image has 3 channels (RGB)
+    if len(image.shape) == 2:  # If grayscale
+        image = tf.stack([image, image, image], axis=-1)
+    elif image.shape[-1] == 4:  # If RGBA
+        image = image[:, :, :3]
+
+    # Resize the image
+    image = tf.image.resize(image, [target_height, target_width])
+
+    # Normalize to [-1, 1]
+    image = tf.cast(image, tf.float32)
+    image = (image / 127.5) - 1
+
+    return image
+
+def process_image(input_image):
+    if input_image is None:
+        return None
+
+    # Get original input image size
+    original_size = input_image.size
+
+    # Preprocess the image
+    processed_input = preprocess_single_image(input_image)
+
+    # Add batch dimension
+    processed_input = tf.expand_dims(processed_input, 0)
+
+    # Generate prediction
+    prediction = generator_g(processed_input)
+
+    # Convert the prediction to displayable format
+    output_image = prediction[0] * 0.5 + 0.5  # Denormalize to [0, 1]
+    output_image = tf.clip_by_value(output_image, 0, 1)
+
+    # Convert to numpy array and then to PIL Image
+    output_array = (output_image.numpy() * 255).astype(np.uint8)
+    output_pil = Image.fromarray(output_array)
+
+    return output_pil
+
+# Create Gradio interface
+demo = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(type="pil", label="Input Image"),
+    outputs=gr.Image(type="pil", label="Generated Output"),
+    title="Image Processing Model",
+    description="Upload an image to see the model's output.",
+)
+
+# Launch the interface
+demo.launch(debug=True)