yolov8模型推理测试代码(pt/onnx)
🦖yolov8训练出来的模型,不使用detect.py代码进行模型测试🦖
pt格式模型测试
import cv2
import os
from ultralytics import YOLO
# 定义输入和输出文件夹路径
input_folder = '/input/folder' # 输入文件夹
output_folder = '/output/folder' # 输出文件夹
# 确认输出文件夹存在
if not os.path.exists(output_folder):os.makedirs(output_folder)
# 加载YOLO模型
model = YOLO('yolov8.pt')
# 遍历输入文件夹中的所有文件
for filename in os.listdir(input_folder):# 仅处理图片文件,确保扩展名为 .jpg, .png 等if filename.endswith(('.jpg', '.png','.bmp')):# 加载图像image_path = os.path.join(input_folder, filename)image = cv2.imread(image_path)if image is None:print(f"图像加载失败: {image_path}")continue# 进行预测results = model(image,device='1')# 获取结果中的标注信息annotated_frame = results[0].plot() # 将检测结果绘制在图像上# 保存结果图像到输出文件夹save_path = os.path.join(output_folder, f'result_{filename}')save_status = cv2.imwrite(save_path, annotated_frame)if save_status:print(f"检测结果保存到: {save_path}")else:print(f"保存图像失败: {save_path}")
onnx模型测试结果
import cv2
import os
import onnxruntime as ort
import numpy as np
# 定义输入和输出文件夹路径
input_folder = '/input/folder' # 输入文件夹
output_folder = '/output/folder' # 输出文件夹
# 确认输出文件夹存在
if not os.path.exists(output_folder):os.makedirs(output_folder)
# 加载 ONNX 模型
onnx_model_path = 'yolov8.onnx'
session = ort.InferenceSession(onnx_model_path, providers=['CUDAExecutionProvider'])
# 设置模型输入的尺寸
input_width, input_height = 640, 640 # 根据模型的输入尺寸
# 遍历输入文件夹中的所有文件
for filename in os.listdir(input_folder):# 仅处理图片文件,确保扩展名为 .jpg, .png 等if filename.endswith(('.jpg', '.png', '.bmp')):# 加载图像image_path = os.path.join(input_folder, filename)image = cv2.imread(image_path)if image is None:print(f"图像加载失败: {image_path}")continue# 调整图像大小并进行预处理input_image = cv2.resize(image, (input_width, input_height))input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)blob = np.expand_dims(input_image.astype(np.float32) / 255.0, axis=0).transpose(0, 3, 1, 2)# 设置输入input_name = session.get_inputs()[0].nameoutputs = session.run(None, {input_name: blob})# 解析预测结果并绘制检测框for detection in outputs[0][0]: # 根据模型的输出格式解析score = float(detection[4]) # 假设第5个值是置信度if score > 0.5: # 置信度阈值,取决于具体任务x_center, y_center, width, height = detection[:4]left = int((x_center - width / 2) * image.shape[1])top = int((y_center - height / 2) * image.shape[0])right = int((x_center + width / 2) * image.shape[1])bottom = int((y_center + height / 2) * image.shape[0])# 绘制检测框cv2.rectangle(image, (left, top), (right, bottom), (0, 255, 0), 2)label = f"{score:.2f}"cv2.putText(image, label, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)# 保存结果图像到输出文件夹save_path = os.path.join(output_folder, f'result_{filename}')save_status = cv2.imwrite(save_path, image)if save_status:print(f"检测结果保存到: {save_path}")else:print(f"保存图像失败: {save_path}")