点击查看代码
import torch
from ultralytics.utils import ASSETS, yaml_load
from ultralytics.utils.checks import check_requirements, check_yaml
import numpy as np
import cv2
import onnxruntime as ort
class YOLOv8:
"""YOLOv8 object detection model class for handling inference and visualization."""
def __init__(self, onnx_model, input_image, confidence_thres, iou_thres):
"""
Initializes an instance of the YOLOv8 class.
Args:
onnx_model: Path to the ONNX model.
input_image: Path to the input image.
confidence_thres: Confidence threshold for filtering detections.
iou_thres: IoU (Intersection over Union) threshold for non-maximum suppression.
"""
self.onnx_model = onnx_model
self.input_image = input_image
self.confidence_thres = confidence_thres
self.iou_thres = iou_thres
self.classes = yaml_load(check_yaml("mycoco.yaml"))["names"]
self.color_palette = np.random.uniform(0, 255, size=(len(self.classes), 3))
def draw_detections(self, img, box, score, class_id):
"""
Draws bounding boxes and labels on the input image based on the detected objects.
Args:
img: The input image to draw detections on.
box: Detected bounding box.
score: Corresponding detection score.
class_id: Class ID for the detected object.
Returns:
None
"""
x1, y1, w, h = box
color = self.color_palette[class_id]
cv2.rectangle(img, (int(x1), int(y1)), (int(x1 + w), int(y1 + h)), color, 2)
label = f"{self.classes[class_id]}: {score:.2f}"
(label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
label_x = x1
label_y = y1 - 10 if y1 - 10 > label_height else y1 + 10
cv2.rectangle(
img, (label_x, label_y - label_height), (label_x + label_width, label_y + label_height), color, cv2.FILLED
)
cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)
x = max(int(x1), 0)
y = max(int(y1), 0)
w = int(x1 + w) - x
h = int(y1 + h) - y
return [x, y, w, h]
def preprocess(self):
"""
Preprocesses the input image before performing inference.
Returns:
image_data: Preprocessed image data ready for inference.
"""
self.img = cv2.imread(self.input_image)
self.img_height, self.img_width = self.img.shape[:2]
img = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (self.input_width, self.input_height))
image_data = np.array(img) / 255.0
image_data = np.transpose(image_data, (2, 0, 1))
image_data = np.expand_dims(image_data, axis=0).astype(np.float32)
return image_data
def postprocess(self, input_image, output):
"""
Performs post-processing on the model's output to extract bounding boxes, scores, and class IDs.
Args:
input_image (numpy.ndarray): The input image.
output (numpy.ndarray): The output of the model.
Returns:
numpy.ndarray: The input image with detections drawn on it.
"""
outputs = np.transpose(np.squeeze(output[0]))
rows = outputs.shape[0]
boxes = []
scores = []
class_ids = []
x_factor = self.img_width / self.input_width
y_factor = self.img_height / self.input_height
for i in range(rows):
classes_scores = outputs[i][4:]
max_score = np.amax(classes_scores)
if max_score >= self.confidence_thres:
class_id = np.argmax(classes_scores)
x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]
left = int((x - w / 2) * x_factor)
top = int((y - h / 2) * y_factor)
width = int(w * x_factor)
height = int(h * y_factor)
class_ids.append(class_id)
scores.append(max_score)
boxes.append([left, top, width, height])
indices = cv2.dnn.NMSBoxes(boxes, scores, self.confidence_thres, self.iou_thres)
for i in indices:
box = boxes[i]
score = scores[i]
class_id = class_ids[i]
rectangle_params = self.draw_detections(input_image, box, score, class_id)
return input_image, rectangle_params
def main(self):
"""
Performs inference using an ONNX model and returns the output image with drawn detections.
Returns:
output_img: The output image with drawn detections.
"""
session = ort.InferenceSession(self.onnx_model, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
model_inputs = session.get_inputs()
input_shape = model_inputs[0].shape
self.input_width = input_shape[2]
self.input_height = input_shape[3]
img_data = self.preprocess()
outputs = session.run(None, {model_inputs[0].name: img_data})
return self.postprocess(self.img, outputs)
class YOLOv8Inference:
def __init__(self, model_path, conf_thres, iou_thres):
check_requirements("onnxruntime-gpu" if torch.cuda.is_available() else "onnxruntime")
self.detection = YOLOv8(model_path, None, conf_thres, iou_thres)
def process_image(self, img_path):
print("207:img_path",img_path)
self.detection.input_image = img_path
output_image, rectangles = self.detection.main()
return output_image, rectangles
点击查看代码
import argparse
import time
import cv2
import datetime
from examples.YOLOV8 import YOLOv8Inference
if __name__ == "__main__":
img_path = "D:/work/wugong/ultralytics-main/TESTONNX/INPUTimg/11.png"
current_time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str,
default="./runs/detect/train12/weights/best.onnx",
help="Input your ONNX model.")
parser.add_argument("--img", type=str, default=img_path, help="Path to input image.")
parser.add_argument("--conf-thres", type=float, default=0.5, help="Confidence threshold")
parser.add_argument("--iou-thres", type=float, default=0.5, help="NMS IoU threshold")
args = parser.parse_args()
yolo_inference = YOLOv8Inference(args.model, args.conf_thres, args.iou_thres)
start = time.time()
output_image, rectangles = yolo_inference.process_image(args.img)
end = time.time()
print("38rectangles:",rectangles)
print(f"total cost time:{end - start}")
cv2.imwrite(f"./TESTONNX/OUTPUTimg/{current_time}.png", output_image)
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】凌霞软件回馈社区,博客园 & 1Panel & Halo 联合会员上线
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】博客园社区专享云产品让利特惠,阿里云新客6.5折上折
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· DeepSeek “源神”启动!「GitHub 热点速览」
· 我与微信审核的“相爱相杀”看个人小程序副业
· 微软正式发布.NET 10 Preview 1:开启下一代开发框架新篇章
· 如何使用 Uni-app 实现视频聊天(源码,支持安卓、iOS)
· C# 集成 DeepSeek 模型实现 AI 私有化(本地部署与 API 调用教程)