from tensorflow.keras.models import load_model
from PIL import Image
import gradio as gr
import numpy as np

model=load_model('NumberToImage.h5')

def generate_image(number):
    number=int(number) 
    if not (0<=number<10):
        return "Number must be between 0 and 9"

    one_hot=np.zeros((1, 10))
    one_hot[0, number]=1
    
    img=model.predict(one_hot)[0]

    if img.max() <= 1.0:
        img=(img * 255).astype(np.uint8)
    else:
        img=img.astype(np.uint8)
    
    img=Image.fromarray(img)
    img=img.resize((280, 280), resample=Image.NEAREST)

    return img

gr.Interface(
    fn=generate_image,
    inputs=gr.Number(label="Enter a number (0 to 9)", precision=0),
    outputs=gr.Image(type="pil"),
    title="Number to Image Generator",
    description="Enter a digit (0 to 9) to generate its corresponding image from a trained model. [Github Repository](https://github.com/NishantkSingh0/Image-Generation/)"
).launch()