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("rushic24/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


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,
                     )
iface.launch(debug=True)