import gradio as gr
import numpy as np
from huggingface_hub import from_pretrained_keras

model = from_pretrained_keras("mouaff25/tyre_quality_classifier")
labels = ["defective", "good"]

def classify_image(inp: np.ndarray):
  inp = inp.reshape((-1, 224, 224, 3))
  prediction = model.predict(inp).flatten()
  confidences = {
    "defective": 1 - float(prediction[0]),
    "good": float(prediction[0])
  }
  return confidences

demo = gr.Interface(fn=classify_image,
             inputs=gr.Image(shape=(224, 224)),
             outputs=gr.Label(num_top_classes=2),
             examples=["./examples/defective.jpg", "./examples/good.jpg"],
             title="Tire Quality Classifier",
             description="This model can distinguish between good and defective tyres. Upload an image of a tyre to see the model in action!")

demo.launch()