app.py

import gradio as gr
from transformers import ViTImageProcessor, ViTForImageClassification
from PIL import Image
import torch

# --- CHANGEMENT CRITIQUE : Charger VOTRE modèle fine-tuné ---
model_name = "MODLI/vit-fashion-classifier"  # <--- REMPLACER par votre modèle entraîné
processor = ViTImageProcessor.from_pretrained(model_name)
model = ViTForImageClassification.from_pretrained(model_name)

# Fonction de prédiction avec seuil de confiance
def predict(image):
    # Pré-traiter l'image exactement comme pendant l'entraînement
    inputs = processor(images=image, return_tensors="pt")
    
    # Prédire
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
    
    # Appliquer softmax pour obtenir les probabilités
    probabilities = torch.nn.functional.softmax(logits, dim=-1)[0]
    top_probs, top_indices = torch.topk(probabilities, 5)  # Top 5 predictions
    
    # Formater les résultats
    predictions = []
    for i, (prob, idx) in enumerate(zip(top_probs, top_indices)):
        pred_label = model.config.id2label[idx.item()]
        confidence = prob.item()
        # N'afficher que si la confiance est > 5%
        if confidence > 0.05 or i == 0:  # Toujours afficher la première même si faible
            predictions.append((pred_label, f"{confidence:.2%}"))
    
    return predictions

# Interface Gradio
title = "Fashion Item Classifier"
description = "Upload an image of a clothing item, and I will classify it."

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Clothing Item"),
    outputs=gr.Label(label="Predictions", num_top_classes=5),
    title=title,
    description=description,
    examples=[["path_to_example_image_1.jpg"], ["path_to_example_image_2.jpg"]], # Ajoutez des exemples
)

demo.launch(debug=True)