Update app.py

app.py

@@ -1,18 +1,50 @@
 import gradio as gr
-import cv2
 import torch
 from PIL import Image
 import numpy as np
+from scipy.ndimage import label, find_objects
 
 # Cargar el modelo YOLO (usando YOLOv5 como ejemplo)
 model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # Puedes cambiar 'yolov5s' por cualquier otro modelo
 
-# Función para realizar detección de objetos
-def detect_objects(image):
-    image = np.array(image)
-    results = model(image)
-    results_image = results.render()[0]
-    return Image.fromarray(results_image)
+# Función para generar la máscara en lugar de la detección de objetos
+def generate_mask(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
 
 # Interfaz de Gradio para cargar una imagen
 def gradio_interface():
@@ -38,7 +70,6 @@ def gradio_interface():
         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 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>")
@@ -46,11 +77,11 @@ def gradio_interface():
         # Organizar imágenes de carga y procesadas en la misma fila
         with gr.Row():
             img_input = gr.Image(label="Upload Image")
-            img_output = gr.Image(label="Image with Detected Objects")
+            img_output = gr.Image(label="Image with Generated Mask")
         
-        # Botón pequeño para clasificación
-        btn_classify = gr.Button("Classify Fermentation Level")
-        btn_classify.click(detect_objects, inputs=img_input, outputs=img_output)
+        # Botón pequeño para clasificación de máscara
+        btn_classify = gr.Button("Generate Mask for Fermentation Level")
+        btn_classify.click(generate_mask, inputs=img_input, outputs=img_output)
 
         # Descripciones de clases de cacao según NTC1252:2021
         gr.Markdown("""