import gradio as gr
import tensorflow as tf
from PIL import Image
import numpy as np
from fastai.vision.all import *


learn = load_learner('export.pkl')
model_path = "keras_model.h5"


model = tf.keras.models.load_model(model_path)
categories = ('Blouse', 'Dress', 'Pants', 'Shirt', 'Shorts')

title = "Clothing Identifier"

class_labels = [
'Cotton',
'Linen',
'Silk',
'Wool',
'Polyester',
'Nylon',
'Rayon',
'Fleece',
'Leather',
'Synth Leather'
]


def classify_image(img):
    pred, idx, probs = learn.predict(img)
    
    img = Image.fromarray((img * 255).astype(np.uint8))
    img = img.resize((224, 224))
    img_array = tf.keras.preprocessing.image.img_to_array(img)
    img_array = tf.expand_dims(img_array, 0)  

    
    predictions = model.predict(img_array)
    predicted_class = class_labels[np.argmax(predictions)]
    highest_prob_index = probs.argmax()

    return { "Material Type is":predicted_class,  
            
            "Cloth Category is ":categories[highest_prob_index]
           }


iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(),
    outputs=gr.Textbox(),
    title = title,
    examples = ['dress.jpg', 'shirt.jpg', 'pants.jpg', 'shorts.jpg'],
    # live=True,
)


iface.launch(share=True)