diff --git a/pt2onnx.py b/pt2onnx.py new file mode 100644 index 00000000..7b109495 --- /dev/null +++ b/pt2onnx.py @@ -0,0 +1,67 @@ +from ultralytics import YOLO +import onnx +import onnxruntime as ort +import yaml +import cv2 +import numpy as np +import glob +import os + +# 1. 加载并导出模型 +model = YOLO('runs/detect/train/weights/best.pt') +model.to('cuda') + +model.export( + format="onnx", # ONNX格式 + dynamic=True, + simplify=True +) + +# 2. 加载ONNX模型 +onnx_model = onnx.load('runs/detect/train/weights/best.onnx') +onnx.checker.check_model(onnx_model) +ort_session = ort.InferenceSession('runs/detect/train/weights/best.onnx') + +# 3. 加载验证集路径 +with open('data.yaml', 'r') as f: + data_config = yaml.safe_load(f) +val_dir = os.path.dirname(data_config['val']) +val_dir = val_dir.replace('../', '') # 直接删除"../"前缀 +print("val_dir:", val_dir) # 调试信息 +val_images = os.path.join(val_dir, 'images') if 'valid' in val_dir else os.path.join(val_dir, 'val', 'images') +val_images = val_images.replace('\\', '/') # 统一使用正斜杠路径分隔符 +print(f"val_images: {val_images}") # 调试信息 + +# 4. 验证集测试 +correct = 0 +total = 0 +print("开始验证...") + +matched_files = glob.glob(val_images + '/*.jpg') +print(f"找到 {len(matched_files)} 张测试图片") # 文件匹配调试 + +for img_path in matched_files[:20]: # 测试前20张 + # 预处理 + img = cv2.imread(img_path) + img = cv2.resize(img, (640, 640)) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 确保RGB顺序 + img = img.transpose(2, 0, 1).astype(np.float32) / 255.0 + img = np.expand_dims(img, axis=0).astype(np.float32) + + # 推理 + outputs = ort_session.run(None, {"images": img}) + + # 添加输出调试信息 + print(f"输出形状: {outputs[0].shape}") + print(f"输出示例数据: {outputs[0][0][:5]}...") + + # 使用置信度阈值判断检测 + confidences = outputs[0][0, 4, :] + print(f"最大置信度: {np.max(confidences):.3f}") + + # 使用阈值判断是否检测到目标 + if np.any(confidences > 0.5): + correct += 1 + total += 1 + +print(f"验证结果: {correct}/{total} 张图片检测到目标") \ No newline at end of file diff --git a/runs/detect/train/weights/best.onnx b/runs/detect/train/weights/best.onnx new file mode 100644 index 00000000..3e21b710 Binary files /dev/null and b/runs/detect/train/weights/best.onnx differ