ComputerVision / src /yolo.py
pirahansiah's picture
add YOLO
317aea0
import cv2
import numpy as np
import gradio as gr
from ultralytics import YOLO
def draw_boxes(image, boxes):
for box in boxes:
x1, y1, x2, y2, name, prob = box
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(image, f"{name} {prob:.2f}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 1.9, (255,55,0), 2)
return image
def detect_objects_on_image(buf):
model = YOLO("files/yolov8n.pt")
results = model.predict(buf)
result = results[0]
output = []
for box in result.boxes:
x1, y1, x2, y2 = [
round(x) for x in box.xyxy[0].tolist()
]
class_id = box.cls[0].item()
prob = round(box.conf[0].item(), 2)
output.append([
x1, y1, x2, y2, result.names[class_id], prob
])
return output
path = [['files/a.png'],['files/image_0.png'],['files/huggingface.png']]
inputs = [
gr.inputs.Image(type="filepath", label="Input Image"),
gr.inputs.Radio(label="YOLO Methods",
choices=[
"v8"
]),
]
outputs = [
gr.outputs.Image(type="numpy", label="Output YOLO Image")
]
def process_image(input_image, radio_choice,slider_val):
img = cv2.imread(input_image)
img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
boxes=detect_objects_on_image(img)
img_with_boxes = draw_boxes(img, boxes)
return img_with_boxes
def on_change(radio_choice):
outputs[0].update(process_image(
inputs[0].value,
radio_choice)
)
yolo = gr.Interface(
fn=process_image,
inputs=inputs,
outputs=outputs,
on_change=on_change,
examples=path,
title="YOLO: Computer Vision and Deep Learning by Farshid PirahanSiah",
)