import gradio as gr
from tensorflow.keras.models import load_model
from PIL import Image
import numpy as np

# Class names mapping
class_names = {
    0: 'arduino',
    1: 'battery',
    2: 'Bluetooth module',  # folder '10'
    3: 'DCmotor',
    4: 'DHT-11',
    5: 'ESP8266',
    6: 'LCD',
    7: 'Loadcell',
    8: 'RFID',
    9: 'Tiva',
    10: 'Ultrasonic',
}


# Load the pre-trained model
model = load_model("electronic_components_classifier_97.keras")

# Function to predict image
def predict_image(img):
    img = img.convert("RGB")
    img = img.resize((224, 224))
    data = np.asarray(img)
    data = data / 255.0
    probs = model.predict(np.expand_dims(data, axis=0))
    top_prob = probs.max()
    top_pred = class_names[np.argmax(probs)]
    return f"This is a {top_pred} with {top_prob * 100:.2f}% confidence."

# Create the Gradio interface
interface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Electronic Component Detector",
    description="Upload an image of an electronic component, and the model will classify it.",
)

# Run the Gradio app
if __name__ == "__main__":
    interface.launch()