Upload 4 files

app.py

@@ -0,0 +1,96 @@
+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((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:
+        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.")

best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:380f9383305dce76f36c8f552ce94e2518f3e7ab1b6aafbe4754935f55a2aeeb
+size 136704041

database_connect.py

@@ -0,0 +1,23 @@
+import os
+import pymssql
+
+
+class DatabaseConnection:
+    def __init__(self, server, database, username, password):
+        self.connection = pymssql.connect(server, username, password, database)
+
+    def insert_data(self, table_name, data_values):
+        columnas = ('img_size', 'intersect_over_union', 'probabilidad', 'user_text', 'feedback_text', 'check_status')
+        insert_query = f"INSERT INTO {table_name} ({','.join([columna for columna in columnas])}) VALUES ({','.join(['%s' for _ in data_values])})"
+        print(insert_query)
+        cursor = self.connection.cursor()
+        cursor.execute(insert_query, data_values)
+        self.connection.commit()
+        cursor.close()
+
+server = os.environ.get('SERVER') # Reemplaza con el nombre de tu servidor
+database = os.environ.get('DATABASE') # Reemplaza con el nombre de tu base de datos
+username = os.environ.get('USERNAME') # Reemplaza con tu nombre de usuario
+password = os.environ.get('PASSWORD') # Reemplaza con tu contraseña
+
+db_connection = DatabaseConnection(server, database, username, password)

requirements.txt

@@ -0,0 +1,19 @@
+streamlit==1.33.0
+ultralytics
+ultralyticsplus
+gradio
+IPython
+matplotlib>=3.2.2
+numpy>=1.18.5
+opencv-python-headless
+pillow
+psutil
+PyYAML>=5.3.1
+scipy>=1.4.1
+torch>=1.7.0
+torchvision>=0.8.1
+tqdm>=4.41.0
+pandas
+thop  # FLOPs computation
+cython==0.29.35
+pymssql==2.2.7