import gradio as gr
from transformers import ConvNextForImageClassification
from PIL import Image
from torchvision import transforms
import torch

# Charger le modèle depuis Hugging Face Hub
model_name = "Muselion/Test_Dogs"  # Mets ton repo ici
model = ConvNextForImageClassification.from_pretrained(model_name)
model.eval()

# Transforms à appliquer à l'image uploadée
val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

def predict(image: Image.Image):
    if image.mode != "RGB":
        image = image.convert("RGB")
    input_tensor = val_transform(image).unsqueeze(0)
    with torch.no_grad():
        outputs = model(input_tensor)
        predicted_class = outputs.logits.argmax(dim=-1).item()

        id2label = model.config.id2label
        label_name = id2label.get(predicted_class, f"Classe({predicted_class})")
    return label_name

# --- Images pour la galerie transformations ---
images_paths = [
    "transforms/original.jpg",
    "transforms/mirror.jpg",
    "transforms/zoom.jpg",
    "transforms/blur.jpg",
    "transforms/rotate.jpg"
]

images = [Image.open(path).convert("RGB") for path in images_paths]

# --- Interface Gradio avec deux onglets ---
with gr.Blocks(css="""
#label-output .gr-text {
    font-size: 24px !important;
    font-weight: bold;
    color: #222;
}
""") as demo:
    with gr.Tab("Classification"):
        gr.Markdown("### Upload une photo de chien et prédire la race")
        image_input = gr.Image(type="pil")
        label_output = gr.Text()
        image_input.change(fn=predict, inputs=image_input, outputs=label_output)

    with gr.Tab("Données d'entrainement"):
        gr.Markdown("### Variations d'une même image pendant l'entraînement")
        gallery = gr.Gallery(value=images, label="Images transformées")

        gr.Markdown("---")
        gr.Markdown("### Répartition des races de chiens dans le jeu d'entraînement")
        gr.Image("transforms/class_distribution.png", label="Répartition des classes")

demo.launch()