import gradio as gr
# import torch
import tensorflow as tf
import numpy as np
from PIL import Image
# from torchvision.transforms import Compose, Resize, ToTensor, Normalize



print('loading model..')
model = tf.keras.models.load_model('model.keras')
print('loaded.')


# transform = Compose([
#     Resize((300,300)),
#     ToTensor(),
#     Normalize(mean=[0.485, 0.456, 0.406],
#                         std=[0.229, 0.224, 0.225]),
# ])

def predict(img):
    if isinstance(img, np.ndarray):
        img = Image.fromarray(img.astype('uint8'), 'RGB')
    img = img.resize((300,300))
    img = np.array(img)
    img = np.expand_dims(img,axis=0)
    labels = list(range(10))
    # img = transform(img)
    # img = img.unsqueeze(0)

    # with torch.inference_mode():
    #     prediction = torch.softmax(model(img),dim=1)[0]
    prediction = model(img)[0]


    result = { num:float(prob.numpy()) for num, prob in enumerate(prediction)}

    
    return result


example_images = [
    "1.png",  # Make sure these paths are correct
    "2.png",
    "3.png"
]


demo = gr.Interface(
    fn=predict,
    inputs=["image"],
    outputs=[gr.Label(num_top_classes=3, label="Predictions")],
    examples=example_images,    
)

demo.launch()