from ultralytics import YOLO
import matplotlib.pyplot as plt
import gradio as gr
import cv2
import numpy as np
import random
import seaborn as sns
from PIL import Image
import io
def predict(path:str):
    model = YOLO("yolov8n.yaml")
    model = YOLO("best.pt")
    image = cv2.imread(path)
    
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    results = model.predict(source=path) 
    paleta= sns.color_palette("bright", 17)
    fig = plt.figure()
    vectObjs = results[0].masks.xy
    resNumClase = results[0].boxes.cls.numpy().astype(int)
    conf = results[0].boxes.conf.numpy()
    for i in range(len(vectObjs)):
        objDet = vectObjs[i].astype(int)
        color = (paleta[i][0]*255, paleta[i][1]*255, paleta[i][2]*255)
        image = cv2.polylines(image, [objDet], True, color, 4)
        plt.text(objDet[0][0], objDet[0][1], results[0].names[resNumClase[i]]+" "+ str(conf[i]), bbox=dict(facecolor=paleta[i], alpha=0.5))
    
    plt.imshow(image)
    plt.axis('off')
    return plt
gr.Interface(fn=predict, 
             inputs=gr.components.Image(type="filepath", label="Input"),
             outputs=gr.Plot(label="Resultado de detección de objetos con regularizacion")).launch()
             #outputs=gr.components.Image(type="pil", label="Output")).launch()