import tensorflow as tf
import numpy as np
import gradio as gr
from PIL import Image

# 1. Charger le modèle entraîné (.h5)
model = tf.keras.models.load_model("CNN_model.h5")  # change le nom si besoin

# 2. Définir les classes CIFAR-10
classes = [
    "airplane", "automobile", "bird", "cat", "deer",
    "dog", "frog", "horse", "ship", "truck"
]

# 3. Fonction de prétraitement + prédiction
def predict(image):
    image = image.resize((32, 32))         # Redimensionner à 32x32
    image_array = np.array(image) / 255.0  # Normaliser
    image_array = image_array.reshape(1, 32, 32, 3)  # Ajouter batch dimension
    predictions = model.predict(image_array)[0]
    result = {classes[i]: float(predictions[i]) for i in range(10)}
    return result

# 4. Interface Gradio
gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="Classification d’Images avec CIFAR-10",
    theme='JohnSmith9982/small_and_pretty',
    description="""Ce projet met en œuvre une application de classification d images basée sur un modèle 
    d apprentissage profond entraîné sur le célèbre jeu de données CIFAR-10, qui comprend 10 catégories 
    d objets (avion, automobile, oiseau, chat, cerf, chien, grenouille, cheval, bateau, camion)"""
).launch()