yolov8 onnxruntime
时间: 2025-01-14 13:57:21 浏览: 49
### 使用ONNX Runtime运行YOLOv8模型
为了成功使用ONNX Runtime来执行YOLOv8模型,需确保所使用的CUDA和cuDNN版本与ONNX Runtime兼容[^1]。以下是具体方法:
#### 安装必要的库
首先安装`onnxruntime-gpu`包以及YOLOv8所需的依赖项。
```bash
pip install onnxruntime-gpu yolov8-torch
```
确认已安装的`cudnn`版本与ONNX Runtime相匹配非常重要。如果当前环境中的cuDNN版本不支持,则可能需要调整环境配置以适应所需的主要版本。
#### 导入必需模块并加载模型
编写Python脚本来导入必要模块,并加载训练好的YOLOv8模型转换成的`.onnx`文件。
```python
import numpy as np
import cv2
import onnxruntime as ort
from pathlib import Path
model_path = "path/to/yolov8.onnx"
session = ort.InferenceSession(model_path)
def preprocess(image):
image_resized = cv2.resize(image, (640, 640))
img_in = cv2.cvtColor(image_resized, cv2.COLOR_BGR2RGB)
img_in = np.transpose(img_in, (2, 0, 1)).astype(np.float32) # HWC -> CHW
img_in /= 255.0
img_in = np.expand_dims(img_in, axis=0)
return img_in
```
这段代码定义了一个预处理函数用于图像输入前准备,包括尺寸调整、颜色空间变换及标准化操作。
#### 执行推理过程
通过调用`InferenceSession.run()`来进行实际预测工作。
```python
image = cv2.imread('test_image.jpg')
input_tensor = preprocess(image)
outputs = session.run(None, {"images": input_tensor})
```
这里假设模型接受名为`"images"`的张量作为输入参数名;对于不同版本或自定义导出设置下的模型,此名称可能会有所不同,请参照具体的文档说明。
#### 后处理检测结果
最后一步是对输出数据进行解析得到最终的目标框位置和其他属性信息。
```python
boxes = outputs[0][..., :4]
scores = outputs[0][..., 4:]
class_ids = scores.argmax(axis=-1)
confidences = np.max(scores, axis=-1)
for box, confidence, class_id in zip(boxes[class_ids != -1], confidences[class_ids != -1], class_ids[class_ids != -1]):
if confidence >= 0.5:
x_center, y_center, width, height = box * [image.shape[1], image.shape[0], image.shape[1], image.shape[0]]
top_left_x = int(x_center - width / 2.)
top_left_y = int(y_center - height / 2.)
bottom_right_x = int(top_left_x + width)
bottom_right_y = int(top_left_y + height)
label = f'{confidence:.2f}'
color = tuple(map(int, np.random.randint(0, 255, size=(3,))))
cv2.rectangle(image, (top_left_x, top_left_y), (bottom_right_x, bottom_right_y), color=color, thickness=2)
cv2.putText(image, label, (top_left_x, top_left_y - 7), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
fontScale=0.5, color=color, thickness=2)
cv2.imshow("Detection", image)
cv2.waitKey()
```
上述代码片段展示了如何绘制边界框并将置信度显示出来。
阅读全文
相关推荐


















