app.py

import gradio as gr
import cv2
import numpy as np
import torch
from PIL import Image
from GAN.infer import ImageColorizationPipeline


pipeline = ImageColorizationPipeline(
    model_path="GAN/pytorch_model.bin",  
    input_size=256,
    model_size="tiny"
)

def colorize_image(image):
    image = image.convert("RGB")
    img_np = np.array(image)
    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
    result_bgr = pipeline.process(img_bgr)
    result_rgb = cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB)
    output_image = Image.fromarray(result_rgb)
    return (image,output_image)

examples = [
        ['./images/image_01.png'],
        ['./images/image_02.png'],
        ['./images/image_03.png'],
        ['./images/image_04.png'],
        ['./images/image_05.png'],
        ['./images/image_06.png'],
        ['./images/image_07.png'],
        ]
demo = gr.Interface(
    fn=colorize_image,
    theme = gr.themes.Soft(),
    inputs=gr.Image(type="pil", label="Upload Grayscale Image", height = 500),
    outputs=gr.ImageSlider(type="pil", label="Colorized Image",height = 500),
    examples=examples,
    title="Grayscale to RGB Image",
    description="Upload a grayscale image. Models generate Color image"
)

demo.launch()