app.py

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


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((int(x * g) for x in im.size), Image.LANCZOS) # resize with antialiasing

    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

# 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)

    else:
        st.warning("Error procesando la imagen. Volver a probar.")