|
import gradio as gr |
|
from ultralytics import YOLO |
|
import cv2 |
|
import numpy as np |
|
import supervision as sv |
|
import os |
|
from pathlib import Path |
|
from supervision.draw.color import Color |
|
|
|
|
|
|
|
title = "Excavator Detector" |
|
description = "An Yolov8(m) object detection computer vision model to detect images of excavators." |
|
|
|
excavator_path = Path('.') |
|
example_path = excavator_path / "examples" |
|
|
|
|
|
example_list = [["examples/" + example] for example in os.listdir(example_path)] |
|
|
|
def predict(img, conf_thres=0.35, iou_thres=0.5): |
|
model = YOLO('./weights/excavator_only_best_v8m.pt') |
|
classNames = ['excavator'] |
|
|
|
results = model.predict(source=img, conf=conf_thres, iou=iou_thres)[0] |
|
detections = sv.Detections.from_yolov8(results) |
|
|
|
box_annotator = sv.BoxAnnotator() |
|
box_annotator.text_scale = 1.3 |
|
box_annotator.text_color = Color.white() |
|
box_annotator.text_thickness = 2 |
|
|
|
labels = [ |
|
f"{classNames[class_id]} {confidence:.2f}" |
|
for class_id, confidence |
|
in zip(detections.class_id, detections.confidence) |
|
] |
|
image = cv2.cvtColor(cv2.imread(img),cv2.COLOR_BGR2RGB) |
|
annotated_image = box_annotator.annotate(image, detections=detections, labels=labels) |
|
|
|
return annotated_image, results.speed |
|
|
|
|
|
|
|
|
|
demo = gr.Interface(fn=predict, |
|
inputs=gr.Image(type="filepath"), |
|
outputs=[gr.Image(type='numpy'), |
|
gr.Text(label="Prediction time (ms)")], |
|
examples=example_list, |
|
title=title, |
|
description=description) |
|
|
|
|
|
demo.launch(debug=False |
|
) |