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

# Cargar el modelo YOLO directamente desde el archivo local
model = YOLO('./best.pt')

# Función para procesar la imagen con YOLO
def detectar_objetos(image):
    # Convertir la imagen de entrada (Gradio) a formato OpenCV (BGR)
    image_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

    # Realiza la detección con YOLO directamente en la imagen
    results = model.predict(source=image_bgr, conf=0.25, save=False)

    # Dibuja las detecciones en la imagen
    annotated_image = image_bgr.copy()
    for box in results[0].boxes:
        # Extraer coordenadas y clase
        x1, y1, x2, y2 = map(int, box.xyxy[0])
        cls = int(box.cls)

        # Colores personalizados para las clases
        if cls == 0:  # Bueno
            color = (0, 255, 0)  # Verde
            label = "Well"
        elif cls == 1:  # 
            color = (0, 0, 255)  # Rojo
            label = "Partially"
        elif cls == 2:  # 
            color = (255, 0, 0)  # Azul
            label = "bad"
        else:
            color = (255, 255, 255)
            label = f"Class {cls}"

        # Dibujar el cuadro y la etiqueta
        cv2.rectangle(annotated_image, (x1, y1), (x2, y2), color, thickness=2)
        cv2.putText(
            annotated_image, label, (x1, y1 - 10),
            cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, thickness=1
        )

    # Convertir la imagen anotada de BGR a RGB para mostrar en Gradio
    annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)

    return annotated_image_rgb

# Crear la interfaz Gradio con ejemplos y diseño personalizado
examples = [["cacao_1.png"], ["cacao_2.jpg"]]  # Asegúrate de que estos archivos estén en la misma carpeta que el script

iface = gr.Interface(
    fn=detectar_objetos,
    inputs=gr.Image(label="Upload Image"),
    outputs=gr.Image(label="Image with Detected Objects"),
    title="Fermentation Level Classification for Cocoa Beans",
    examples=examples,
    description="""
        **Cacao Classes According to NTC1252:2021:**
        - **a) Well-fermented:** Optimal fermentation process.
        - **b) Partially fermented:** Incomplete fermentation process.
        - **c) Non-fermented:** Lack of adequate fermentation.
        
        **Explanation According to NTC1252:2021:** Here you can explain how the NTC1252:2021 norm applies to the classification of fermentation levels.
    """
)

# Ejecutar la aplicación
if __name__ == "__main__":
    iface.launch()