Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
-from huggingface_hub import hf_hub_download
 import torch
+from huggingface_hub import hf_hub_download
 from PIL import Image
 from torchvision import transforms
 from skimage.color import rgb2lab, lab2rgb
@@ -13,7 +13,6 @@ repo_id = "Hammad712/GAN-Colorization-Model"
 model_filename = "generator.pt"
 model_path = hf_hub_download(repo_id=repo_id, filename=model_filename)
 
-# Define the generator model (same architecture as used during training)
 from fastai.vision.learner import create_body
 from torchvision.models import resnet34
 from fastai.vision.models.unet import DynamicUnet
@@ -52,21 +51,23 @@ def colorize_image(img, model):
     for img in Lab:
         img_rgb = lab2rgb(img)
         rgb_imgs.append(img_rgb)
-    return np.stack(rgb_imgs, axis=0)[0]  # Return the first (and only) image
+    return np.stack(rgb_imgs, axis=0)
 
-# Gradio interface function
-def colorization_function(image):
-    colorized_image = colorize_image(image, G_net)
-    return colorized_image
+# Gradio interface
+def colorize(img):
+    colorized_images = colorize_image(img, G_net)
+    colorized_image = colorized_images[0]
+    return Image.fromarray((colorized_image * 255).astype(np.uint8))
 
-# Gradio Interface Setup
+# Create the Gradio interface
 iface = gr.Interface(
-    fn=colorization_function,
-    inputs=gr.inputs.Image(type="pil", label="Upload Grayscale Image"),
-    outputs=gr.outputs.Image(type="numpy", label="Colorized Image"),
-    title="Image Colorization with GAN",
-    description="Upload a grayscale image, and the model will colorize it using AI."
+    fn=colorize,
+    inputs=gr.Image(type="pil", label="Upload Grayscale Image"),
+    outputs=gr.Image(type="pil", label="Colorized Image"),
+    title="AI Image Colorization",
+    description="Upload a black and white image, and the AI will colorize it.",
+    allow_flagging="never"
 )
 
-# Launch the Gradio interface
+# Launch the interface
 iface.launch()