deploy-gradio / app.py
fcernafukuzaki's picture
Create app.py
14741b7 verified
raw
history blame
2.15 kB
import os
import pandas as pd
import torch
from PIL import Image
from ultralytics import YOLO
import gradio as gr
class YOLODetect():
def __init__(self, modelo):
self.modelo = modelo
def predecir(self, source, imgsz=1280, conf=0.7, iou=0.50):
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
# Inicializa el modelo YOLOv8
path_best_model = 'yolov8n.pt'
modelo_yolo = YOLO(path_best_model)
def detect_objects(im, size, iou, conf):
'''Wrapper para 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.
return "Se guard贸 el feedback exitosamente."
except Exception as err:
print(err)
return "Error al guardar el feedback."
# Configura la interfaz de Gradio
with gr.Blocks() as demo:
gr.Markdown("# YOLOv8 Detecci贸n de objetos")
with gr.Row():
iou_threshold = gr.Slider(label="NMS IoU Threshold (0.0 - 1.0)", minimum=0.0, maximum=1.0, value=0.8)
conf_threshold = gr.Slider(label="Umbral o threshold (0.0 - 1.0)", minimum=0.0, maximum=1.0, value=0.9)
size = gr.Dropdown(label="Tama帽o de la