import gradio as gr
from image_classifier import ImageClassifier
from PIL import Image

classifier = ImageClassifier('models/trained_model.pth')
examples = [
    'examples/0ab9373f-4d97-456a-9b9d-60b4d05d102e.jpg',
    'examples/093e6044-5809-44d6-a93f-f06a6702f20d.jpg',
    'examples/332b50b9-ab27-42ac-9118-ae32b3458d97.jpg',
    'examples/357fbd95-ff50-4d90-9487-531595f02a9e.jpg',
    'examples/b100a944-0e67-4e4d-9ccc-c6c63fb99c8e.jpg'
]

def predict(image):
    image = Image.fromarray(image.astype('uint8'), 'RGB')
    return classifier.classify_image(image)


image = gr.inputs.Image(shape=(192,192))
label = gr.outputs.Label()
intf = gr.Interface(fn=predict, inputs=image, outputs=label, examples=examples)
intf.launch(inline=False)