app.py

import gradio as gr
import pandas as pd
import os
from PIL import Image
from ultralytics import RTDETR
from ultralytics import SAM
import cv2 as cv
import numpy as np
import re

etiquetas_filtradas=[]
cuenta = []

#Modelo de detección de objetos RT-DETR
model = RTDETR('Modelo/rtdetr-l_proyecto.pt')

#Modelo de segmentación de objetos MobileSAM basado en SAM
model_SAM = SAM('Modelo/mobile_sam.pt')

# DataFrame inicial
data = {
    "Comidas": ["Desayuno", "Almuerzo", "Comida", "Merienda", "Cena"],
    "Total": [0, 0, 0, 0, 0]
}
resumen = pd.DataFrame(data)

def obtencion_etiquetas(result):

    object_count = {}

    #Obtención de las etiquetas y el número de veces que se han detectado en la imagen
    for box in result.boxes:
      object_label = result.names[box.cls[0].item()]
      if object_label in object_count:
          object_count[object_label] += 1
      else:
          object_count[object_label] = 1

    global etiquetas_filtradas, cuenta

    etiquetas_filtradas = []
    cuenta = []
        
    for label, count in object_count.items():
      # Obtiene la etiqueta correspondiente y la agrega a la lista
      etiquetas_filtradas.append(label)
      cuenta.append(count)
      print(f"{label}: {count}")      
    print('\n')
    print(etiquetas_filtradas)
    print('\n')
    print(cuenta)

def cargar_dataframe_1(eleccion):

    global etiquetas_filtradas, cuenta

    df = pd.read_csv('Dataset/Tabla_alimentos.csv')
    global nuevo_df, resumen
    
    nombre_primera_columna = df.columns[0]
    columnas_siguientes = df.columns[5:8]
    nuevo_df = df[[nombre_primera_columna] + list(columnas_siguientes)]

    # Filtrar las filas que corresponden a los nombres dados
    nuevo_df = nuevo_df[nuevo_df[nombre_primera_columna].isin(etiquetas_filtradas)]

    # Add a new column "Uds" with the values from contador
    nuevo_df['Uds'] = cuenta

    # Obtener las kcal multiplicadas por la cantidad correspondiente en la columna de Uds
    nuevo_df['Kcal'] = round(nuevo_df['Uds'] * nuevo_df['Kcal'],2)
    
    version2 = cargar_dataframe_2()
    suma = 0.0
    # Utilizar un bucle for para sumar todas las filas de la columna
    for valor in nuevo_df['Kcal']:
        suma += float(valor)
        
    suma = round(suma, 2)
    mod_resumen_dataframe(eleccion, suma)
    
    return nuevo_df, f"{suma} kcal", version2
    
def cargar_dataframe_2():

    global etiquetas_filtradas, cuenta
    
    df = pd.read_csv('Dataset/Tabla_alimentos.csv')
    
    nombre_primera_columna = df.columns[0]
    nombre_segunda_columna = df.columns[6]
    columnas_siguientes = df.columns[8:11]
    nuevo_df = df[[nombre_primera_columna] + [nombre_segunda_columna] + list(columnas_siguientes)]

    # Filtrar las filas que corresponden a los nombres dados
    nuevo_df = nuevo_df[nuevo_df[nombre_primera_columna].isin(etiquetas_filtradas)]
    
    nuevo_df['Uds'] = cuenta

    nuevo_df['Grasas'] = round(nuevo_df['Uds'] * nuevo_df['Grasas'],2)
    nuevo_df['Carbh'] = round(nuevo_df['Uds'] * nuevo_df['Carbh'],2)
    nuevo_df['Prot'] = round(nuevo_df['Uds'] * nuevo_df['Prot'],2)
    
    return nuevo_df

def cargar_dataframe_resumen(eleccion):

    global resumen
    
    suma = 0.0
    # Utilizar un bucle for para sumar todas las filas de la columna
    for valor in resumen['Total']:
        suma += float(valor)
    suma = round(suma, 2)
    return resumen, f"{suma} kcal"

def mod_resumen_dataframe(eleccion, suma):

    global resumen

    if(eleccion == "Desayuno"):
        resumen["Total"][0] = suma
    elif (eleccion == "Almuerzo"):
        resumen["Total"][1] = suma
    elif (eleccion == "Comida"):
        resumen["Total"][2] = suma
    elif (eleccion == "Merienda"):
        resumen["Total"][3] = suma
    elif (eleccion == "Cena"):
        resumen["Total"][4] = suma
    
def umbralizar_imagen(imagen, umbral):

  # Obtener la imagen original
  ancho, alto, canales = imagen.shape

  # Definir la imagen de salida
  imagen_umbralizada = np.zeros((ancho, alto, canales), dtype=np.uint8)

  # Recorrer la imagen, pixel por pixel
  for x in range(ancho):
    for y in range(alto):
      # Obtener el brillo del pixel
      brillo = imagen[x, y]

      # Si todos los canales del pixel tienen un brillo mayor o igual al umbral, copiarlo a la imagen de salida
      if np.any(brillo >= umbral):
        imagen_umbralizada[x, y] = imagen[x, y]

  return imagen_umbralizada
    
def image_mod(Original):
    ancho, alto = Original.size
    Original = Original.resize((640, 640))
    results = model(Original)  # results list
    
    result = results[0]

    Original = np.array(Original)

    #Baja el brillo de la imagen para mejorar la umbralización
    factor_brillo = 0.23
    img_oscura = np.clip(Original * factor_brillo, 0, 255).astype(np.uint8)
    
    for r  in results:
        boxes = r.boxes.xyxy
        
    if boxes is not None and len(boxes) > 0:

        results_SAM = model_SAM(img_oscura, bboxes = boxes)
    
        for x in results_SAM:
            mascara_objeto = x.plot() 
        
        obtencion_etiquetas(result)
        
        # Definir el umbral de brillo
        umbral = 60
    
        # Umbralizar la imagen para extraer solo las máscaras de los objetos
        imagen_umbralizada = umbralizar_imagen(mascara_objeto, umbral)
    
        # Convertir la imagen a escala de grises
        img_gris = cv.cvtColor(imagen_umbralizada, cv.COLOR_BGR2GRAY)
    
        # Aplicar umbralización de Otsu
        _, th = cv.threshold(img_gris, 0, 255, cv.THRESH_BINARY + cv.THRESH_OTSU)
    
        #Bajar el brillo de la imagen para que destaquen más las máscaras de los 
        #objetos segmentados
        factor_brillo = 0.55
        img_mod = np.clip(Original * factor_brillo, 0, 255).astype(np.uint8)
        
        # Hacer la capa blanca más transparente y combinar la imagen con las máscaras
        alpha = 0.4
        img_resultante = cv.addWeighted(img_mod, 1 - alpha, cv.cvtColor(th, cv.COLOR_GRAY2BGR), alpha, 0)
    
    else:
        factor_brillo = 0.55
        img_resultante = np.clip(Original * factor_brillo, 0, 255).astype(np.uint8)
        
    img_resultante = cv.cvtColor(img_resultante, cv.COLOR_BGR2RGB)

    image_with_predictions  = Image.fromarray(img_resultante[..., ::-1])  # RGB PIL image

    image_with_predictions = image_with_predictions.resize((ancho, alto))

    return image_with_predictions 


Imagen = gr.Interface(
    image_mod,
    gr.Image(type="pil"),
    "image",
    examples=[
        #os.path.join(os.path.abspath(''), "Demo_web/1.jpg"),
    ],
    description = "Detección y segmentación de los alimentos de la imagen aportada por el usuario.",
)

Tabla1 = gr.Interface(
    fn= cargar_dataframe_1,
    inputs = gr.Dropdown(
        ["Desayuno", "Almuerzo", "Comida", "Merienda", "Cena"], label="Comida", info="Comida del día"
    ),
    live=True,
    outputs=[gr.Dataframe(
            headers=["Alimento", "Ración(g)", "Uds", "Kcal"],
            label="Resumen alimenticio",
    ), gr.Textbox(label="Total kcal"), 
    gr.Dataframe(
            headers=["Alimento", "Grasas", "Carbh", "Prot"],
            label="Más detalles")],
    description="Desglose de los alimentos y sus calorías, así como recuento total.",
)

Tabla2 = gr.Interface(
    fn=cargar_dataframe_resumen,
    inputs = gr.Textbox(visible=False,label="Actualizar"),
    live=True,
    outputs=[gr.Dataframe(
            resumen,
            label="Resumen alimenticio"
    ), gr.Textbox(label="Total kcal del día")],
    description="Resumen del las kcal ingeridas durante el día.",
)

demo = gr.TabbedInterface([Imagen, Tabla1, Tabla2], ["Inicio", "Detalles", "Diario"], 
                          title = "Detección de alimentos", theme='gradio/soft')

if __name__ == "__main__": 
    demo.launch()