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

iface = gr.Interface(
    fn=generate_mask,
    inputs=gr.Image(label="Upload Image"),
    outputs=gr.Image(label="Image with Generated Mask"),
    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()