Spaces:
Sleeping
Sleeping
File size: 4,028 Bytes
c1a81e4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
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.")
|