app.py

import os
import pandas as pd
import streamlit as st
import torch
from PIL import Image
from ultralytics import YOLO
from database_connect import db_connection


class YOLODetect():
    def __init__(self, modelo):
        self.modelo = modelo

    def predecir(self, source, imgsz=1280, conf=0.7, iou=0.50):
        # conf	float	0.25 umbral de confianza del objeto para la detección
        # iou	float	0.7 umbral de intersección sobre unión (IoU) para NMS
        self.results = self.modelo.predict(source=source, save=True, imgsz=imgsz, conf=conf, iou=iou)
        return self.results

    def render(self):
        result = self.results[0]
        file_name = os.path.join(result.save_dir, result.path)
        render = Image.open(file_name)
        return render

path_best_model = 'best.pt'
modelo_yolo = YOLO(path_best_model)

def detect_objects(size, iou, conf, im):
    '''Wrapper fn for gradio'''
    g = (int(size) / max(im.size))  # gain
    im = im.resize(tuple([int(x * g) for x in im.size]), Image.LANCZOS) # resize with antialiasing
    #im = im.resize((int(x * g) for x in im.size), Image.LANCZOS) # resize with antialiasing
    #im = np.asarray(im, dtype=np.float32)

    model = YOLODetect(modelo_yolo)
    results = model.predecir(source=im, imgsz=int(size), conf=conf, iou=iou)

    objects_detected = results[0].boxes.cls.tolist() # Clases detectadas.
    objects_conf = results[0].boxes.conf.tolist() # Probabilidad de detección por clase detectada.
    
    objects_nested_list = pd.DataFrame({'Clase': objects_detected, 'Probabilidad': objects_conf})

    result_img = model.render()
    return result_img, objects_nested_list

def save_feedback(size, iou, conf, 
                object_count_detected,
                objects_list,
                user_text, feedback_text, check_status):
    try:
        db_connection.insert_data('lego.lego_detectado', (size, iou, conf, user_text, feedback_text, check_status))
        st.success("Se guardó el feeback exitosamente.")
    except Exception as err:
        print(err)
        st.warning("Error al guardar el feedback.")

# Streamlit app layout
st.title('YOLOv8 Detección de figuras LEGO')

# Input
col1, col2 = st.columns(2)
with col1:
    iou_threshold = st.slider("NMS IoU Threshold (0.0 - 1.0)", 0.0, 1.0, 0.8, key="iou")
    conf_threshold = st.slider("Umbral o  threshold (0.0 - 1.0)", 0.0, 1.0, 0.9, key="conf")
with col2:
    size = st.selectbox("Tamaño de la imagen", options=["640", "1280"], key="size")
    uploaded_image = st.file_uploader("Cargar imagen", type=["jpg", "jpeg", "png"], key="image")

# Process uploaded image
if uploaded_image is not None:
    image = Image.open(uploaded_image)
    result_image, objects_nested_list = detect_objects(size=int(size), iou=iou_threshold, conf=conf_threshold, im=image)
    object_count = len(objects_nested_list)

    if result_image is not None:
        col1, col2 = st.columns(2)
        with col1:
            st.image(image, caption="Imagen original", use_column_width=True)
        with col2:
            st.image(result_image, caption="Resultado", use_column_width=True)
        
        with st.form("my_form", clear_on_submit=True):
            st.title("Formulario para feedback")
            st.write(f'Cantidad detectados: {object_count}')
            st.table(objects_nested_list)
            check_status = st.checkbox("¿El resultado contiene la cantidad correcta de figuras detectadas?", value=False)
            user_text = st.text_input("Ingrese el nombre del usuario que realizó la prueba (máximo 50 caracteres)", max_chars=50)
            feedback_text = st.text_input("Ingrese su feedback (máximo 100 caracteres)", max_chars=100)
            # save_button = st.button("Guardar feedback")
            save_button = st.form_submit_button('Guardar feedback')
            if save_button:
                save_feedback(size=int(size), iou=iou_threshold, conf=conf_threshold, 
                            object_count_detected=object_count,
                            objects_list=objects_nested_list,
                            user_text=user_text, feedback_text=feedback_text, check_status=check_status)
    else:
        st.warning("Error procesando la imagen. Volver a probar.")