import gradio as gr
import os
from huggingface_hub import from_pretrained_keras
import tensorflow as tf
import numpy as np

model = from_pretrained_keras("keras-io/bit")
allImages = []
directory = 'images'
CLASSES = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']
# iterate over files in images directory
for filename in os.listdir(directory):
    f = os.path.join(directory, filename)
    if os.path.isfile(f):
      allImages.append(f)

def flower_classifier(image):
  print(image.shape)
  image = tf.image.resize(image, (224, 224))
  image = image / 255.0
  image = tf.expand_dims(image, 0)
  pred = np.argmax(model(image))
  label = CLASSES[pred]
  return label

title = "Image Classification using BigTransfer (BiT)"
description = "This space finetunes BigTransfer (BiT) to classify images from Flower dataset"
article = """<p style='text-align: center'>
        <a href='https://keras.io/examples/vision/bit/' target='_blank'>Keras Example given by Sayan Nath</a>
        <br>
        Space by @rushic24
    </p>
    """

iface = gr.Interface(flower_classifier,
                     title = "Image Classification using BigTransfer (BiT)",
                     inputs = gr.inputs.Image(),
                     outputs = gr.outputs.Label(num_top_classes=5),
                     capture_session=True,
                     examples = allImages,
                     title=title,
                     description=description,
                     article=article)
iface.launch(debug=True)