import tensorflow as tf
import numpy as np
import gradio as gr
from tensorflow.keras.applications.efficientnet import preprocess_input
from tensorflow.keras.preprocessing import image as keras_image

# Load models
modelfor_skin = tf.keras.models.load_model("skin_type_model.keras")
modelfor_disease = tf.keras.models.load_model("skin_disease_model.keras")

# Class names
Skin_typeclass_names = ['Dry', 'Normal', 'Oily']
disease_class_names = ['Acne', 'Blackheads', 'Dark Spots', 'Wrinkles', 'Skin Redness', 'pores', 'Eye Bags']

# Prediction function
def predict_all(img):
    img = img.resize((224, 224))
    img_array = keras_image.img_to_array(img)
    img_array = tf.expand_dims(img_array, axis=0)
    img_array = preprocess_input(img_array)

    # Predict Skin Type
    skin_pred = modelfor_skin.predict(img_array, verbose=0)[0]
    skin_index = np.argmax(skin_pred)
    skin_conf = skin_pred[skin_index]
    skin_result = f"{Skin_typeclass_names[skin_index]} ({skin_conf*100:.2f}%)"

    # Predict Skin Condition
    disease_pred = modelfor_disease.predict(img_array, verbose=0)[0]
    disease_index = np.argmax(disease_pred)
    disease_conf = disease_pred[disease_index]
    disease_result = f"{disease_class_names[disease_index]} ({disease_conf*100:.2f}%)"

    # All probabilities
    prob_list = [
        f"{disease_class_names[i]}: {disease_pred[i]*100:.2f}%"
        for i in range(len(disease_class_names))
    ]
    prob_text = "\n".join(prob_list)

    return skin_result, disease_result, prob_text

# Gradio Interface
iface = gr.Interface(
    fn=predict_all,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Textbox(label="Predicted Skin Type"),
        gr.Textbox(label="Predicted Skin Condition"),
        gr.Textbox(label="All Skin Condition Probabilities")
    ],
    title="Skin Type and Skin Condition Predictor",
    description="Upload a facial skin image to get skin type and condition predictions."
)

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