import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
import os

# Load the model
model_path = "best.pt"
model = YOLO(model_path)

def predict(img):
    if img is None:
        return None, "Veuillez uploader une image."
    
    # Run inference
    results = model(img)
    
    # Get annotated image
    # results[0].plot() returns a numpy array (BGR)
    annotated_img = results[0].plot()
    
    # Convert BGR to RGB for Gradio
    annotated_img_rgb = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
    
    # Count detections
    counts = {}
    for box in results[0].boxes:
        cls = int(box.cls[0])
        label = model.names[cls]
        counts[label] = counts.get(label, 0) + 1
    
    # Format stats string
    if not counts:
        stats = "Aucun objet détecté."
    else:
        stats = "### Détections :\n"
        for label, count in counts.items():
            stats += f"- **{label}**: {count}\n"
            
    return annotated_img_rgb, stats

# Create Gradio interface with a premium look
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="slate")) as demo:
    gr.Markdown(
        """
        # 🌊 SATCAP-OCEANS : Détection de déchet plastique en milieu aquatique
        ### Analyse d'images pour la surveillance des océans et des côtes.
        
        Déposez une image ci-dessous pour détecter automatiquement les déchets plastique en milieu aquatique.
        """
    )
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="numpy", label="Image Aquatique")
            btn = gr.Button("Lancer l'Analyse", variant="primary")
            
        with gr.Column():
            output_img = gr.Image(label="Résultat de l'Analyse")
            output_stats = gr.Markdown(label="Statistiques")
            
    btn.click(fn=predict, inputs=input_img, outputs=[output_img, output_stats])
    
    gr.Examples(
        examples=["example1.png", "example2.png"],
        inputs=input_img
    )
    
    gr.Markdown(
        """
        ---
        Propulsé par **CosmoLABHub** | Modèle YOLOv8
        """
    )

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