import gradio as gr
from PIL import Image
import numpy as np
from scipy.ndimage import label, find_objects

# Función para generar la máscara
def generate_mask(image):
    # Convertir la entrada en un objeto PIL.Image si no lo es
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)

    # Convertir la imagen a RGB y array de numpy
    image_rgb = image.convert("RGB")
    np_image = np.array(image_rgb)

    # Aplicar los límites RGB basados en los análisis anteriores
    min_red, max_red = 21, 183
    min_green, max_green = 0, 142
    min_blue, max_blue = 0, 124

    # Generar la máscara binaria
    mask = (
        (np_image[:, :, 0] >= min_red) & (np_image[:, :, 0] <= max_red) &
        (np_image[:, :, 1] >= min_green) & (np_image[:, :, 1] <= max_green) &
        (np_image[:, :, 2] >= min_blue) & (np_image[:, :, 2] <= max_blue)
    ).astype(np.uint8)

    # Invertir la máscara para que los objetos de fondo (valor 0) se conviertan en 1 y se etiqueten
    inverted_mask = 1 - mask
    labeled_mask, num_features = label(inverted_mask)
    
    # Filtrar objetos pequeños de valor 0
    for i, region in enumerate(find_objects(labeled_mask)):
        if region is not None:
            # Calcular el área del objeto con valor 0
            area = np.sum(labeled_mask[region] == i + 1)
            # Si el área es menor a 3000 píxeles, establece el objeto en blanco
            if area < 3000:
                inverted_mask[region] = np.where(labeled_mask[region] == i + 1, 0, inverted_mask[region])

    # Invertir de nuevo para restaurar los valores de la máscara original
    filtered_mask = 1 - inverted_mask

    # Convertir la máscara modificada a una imagen en blanco y negro
    mask_image = Image.fromarray((filtered_mask * 255).astype(np.uint8))

    return mask_image

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

# 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; 
}
"""

# Configuración de la interfaz con elementos adicionales
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 Generated Mask")
    
    # Botón para generar la máscara
    btn_classify = gr.Button("Generate Mask for Fermentation Level")
    btn_classify.click(generate_mask, inputs=img_input, outputs=img_output)

    # Añadir las imágenes de ejemplo
    gr.Examples(
        examples=examples,
        inputs=img_input,
        label="Example Images"
    )

    # Descripción de las clases de cacao
    gr.Markdown("""
    **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.
    """)
    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()