import gradio as gr
import tensorflow as tf
import numpy as np

model = tf.keras.models.load_model('model.h5')
labels = ['Class_A', 'Class_B'] # Update this after training

def predict(image):
    image = tf.image.resize(image, (224, 224))
    image = np.expand_dims(image, axis=0) / 255.0
    prediction = model.predict(image)[0]
    return {labels[i]: float(prediction[i]) for i in range(len(labels))}

interface = gr.Interface(fn=predict, inputs='image', outputs='label')
interface.launch()