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:  # Well
            color = (0, 255, 0)  # Verde
            label = "Well"
        elif cls == 1:  
            color = (255, 0, 0)  # Rojo
            label = "Partially"
        elif cls == 2:  
            color = (0, 0, 255)  # 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=0
        )

    # 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

# CSS personalizado para la interfaz
css = """
.app-container { background-color: white; }
.title-container {
    background-color: white;
    display: flex;
    align-items: center;
    justify-content: center;
    height: 300px;
    font-size: 24px;
    font-weight: bold;
    text-align: center;
}
.centered-image { 
    display: block; 
    margin: auto; 
}
"""

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

with gr.Blocks(css=css) as demo:
    # Logo y título en la parte superior
    with gr.Row():
        gr.Image(value="Cacaotin.png", width=300, height=300, elem_id="centered-image")
        gr.Markdown("<div class='title-container'>Fermentation Level Classification for Cocoa Beans</div>")

    # Botón de GitHub centrado
    gr.Markdown("<center><a href='https://github.com/kebincontreras/cocoa_beans_interfaces' target='_blank'><button style='background-color: #007bff; color: white; padding: 10px 20px; border: none; border-radius: 5px; font-size: 16px;'>View on GitHub</button></a></center>")

    # Fila para cargar la imagen y mostrar la salida procesada
    with gr.Row():
        img_input = gr.Image(label="Upload Image")
        img_output = gr.Image(label="Image with Detected Objects")
    
    # Botón para procesar la imagen
    btn_detectar = gr.Button("Detect Objects")
    btn_detectar.click(detectar_objetos, inputs=img_input, outputs=img_output)

    # Descripción de las clases de cacao
    gr.Markdown("""
    **Cacao Classes According to NTC1252:2021:**
    - **a) Well-fermented:** 
    - **b) Partially fermented:**
    - **c) Non-fermented:**
    """)
    gr.Image(value="cacao.png", label="a) Well-fermented, b) Partially fermented, c) Non-fermented")

    # Explicación final según la norma
    gr.Markdown("**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__":
    demo.launch()