Create app.py

app.py

@@ -0,0 +1,95 @@
+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 = 'yolov8n.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
+
+    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:
+        # Aquí puede ir el código para almacenar los datos en una base de datos.
+        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 objetos')
+
+# 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.")