import gradio as gr
from transformers import AutoModelForImageClassification, AutoImageProcessor
import torch
from PIL import Image

# 1. Charger le modèle et son processeur
# Le modèle est public : il sera automatiquement téléchargé et mis en cache.
# Vous ne le téléchargez pas vous-même.
model = AutoModelForImageClassification.from_pretrained("anismizi/skin-type-classifier")
processor = AutoImageProcessor.from_pretrained("anismizi/skin-type-classifier")

def predict_skin_type(image):
    """Fonction qui prend une image et retourne le type de peau prédit."""
    inputs = processor(images=image, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model(**inputs)
        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
        predicted_class = predictions.argmax().item()
        confidence = predictions[0][predicted_class].item()
    
    labels = ["dry", "oily"]  # Les étiquettes du modèle
    return f"Type de peau prédit : **{labels[predicted_class]}** (confiance : {confidence:.2%})"

# 2. Créer l'interface Gradio
iface = gr.Interface(
    fn=predict_skin_type,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Analyseur de Type de Peau",
    description="Téléchargez une image de votre visage."
)

# 3. Lancer l'application
iface.launch()