import tensorflow as tf
import gradio as gr

CLASS_NAMES = ['Cat', 'Dog', 'Rabbit']

# Load and initialize tf lite model
interpreter = tf.lite.Interpreter(model_path="classifier.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
input_shape = input_details[0]['shape']

def classify_image(image):
    image = tf.image.resize(image, (224, 224))
    input_data = tf.expand_dims(image, axis=0)
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    predictions = interpreter.get_tensor(output_details[0]['index']).ravel()
    confidences = {CLASS_NAMES[i]: float(predictions[i]) for i in range(len(CLASS_NAMES))}
    return confidences

gr.Interface(fn=classify_image, 
             inputs=gr.Image(shape=(224,224)),                                    
             outputs=gr.Label(num_top_classes=3),                              
             examples=['Cat.jpg', 'Dog.jpg', 'Rabbit.jpg'],
             cache_examples=True).launch()