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

model = load_model("scalp_mobilenetv2_model.h5")

# Update class labels based on your training class_indices
class_labels = ['dry', 'healthy', 'oily']  # <-- Double-check this order

def preprocess_image(image):
    image = image.resize((128, 128))
    image = img_to_array(image) / 255.0
    image = np.expand_dims(image, axis=0)
    return image

def predict_scalp(image):
    processed_image = preprocess_image(image)
    prediction = model.predict(processed_image)
    print("Raw prediction:", prediction[0])  # <-- Debug
    class_index = np.argmax(prediction[0])
    confidence = float(np.max(prediction[0]))
    return {class_labels[class_index]: confidence}

interface = gr.Interface(
    fn=predict_scalp,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="Scalp Condition Classifier",
    description="Upload a scalp image (128x128+). Model predicts dry, healthy, or oily."
)

if __name__ == "__main__":
    interface.launch()