testing / app.py
SahilCarterr's picture
Update app.py
f6600f7 verified
raw
history blame contribute delete
527 Bytes
from ultralytics import YOLO
from PIL import Image
import gradio as gr
import numpy as np
def function(image):
model = YOLO('best_yolo_2.pt')
image = np.array(image)
results = model.predict(image)
img=Image.fromarray(results[0].plot())
num_result=len(results[0].boxes.cls)
names_output=[]
for i in range(num_result):
name=results[0].names[int(results[0].boxes.cls[i])]
names_output.append(name)
return names_output,img
demo=gr.Interface(fn=function,inputs="image",outputs=["text","image"])
demo.launch()