import gradio as gr
import tensorflow as tf
import numpy as np
from tensorflow.keras.preprocessing.image import img_to_array, load_img

# Load the trained model
model = tf.keras.models.load_model('hair_model.h5')

# Define labels for classification
labels = ['curly', 'straight', 'kinky', 'wavy', 'dreadlocks']

# Image preprocessing function
def preprocess_image(image, img_height=299, img_width=299):
    image = image.resize((img_height, img_width))
    image = img_to_array(image) / 255.0  # Rescale the image
    return np.expand_dims(image, axis=0)  # Add batch dimension

# Prediction function
def predict_hair_type(image):
    image = preprocess_image(image)
    predictions = model.predict(image)
    predicted_label = labels[np.argmax(predictions)]
    confidence = np.max(predictions)
    return f"{predicted_label} ({confidence:.2%} confidence)"

# Gradio interface
iface = gr.Interface(
    fn=predict_hair_type,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Hair Type Classifier",
    description="Upload an image to predict the hair type (curly, straight, kinky, wavy, or dreadlocks)."
)

# Launch the Gradio app
if __name__ == "__main__":
    iface.launch()