ShaoHua
104 lines
3.4 KiB
Python
104 lines
3.4 KiB
Python
|
||
import os
|
||
import datetime
|
||
from ultralytics import YOLO
|
||
import argparse
|
||
from model_utils import get_project_root, find_latest_model
|
||
|
||
|
||
def export_to_onnx(model_path=None, imgsz=640, simplify=True, opset=12):
|
||
"""将 PyTorch 模型导出为 ONNX 格式
|
||
|
||
Args:
|
||
model_path: 模型路径（如果为 None 则自动查找）
|
||
imgsz: 输入图片尺寸
|
||
simplify: 是否简化模型
|
||
opset: ONNX opset 版本
|
||
|
||
Returns:
|
||
导出的 ONNX 模型路径
|
||
"""
|
||
project_root = get_project_root()
|
||
|
||
# 如果没提供模型路径，自动查找最新模型
|
||
if model_path is None or not os.path.exists(model_path):
|
||
found_model = find_latest_model()
|
||
if found_model:
|
||
model_path = found_model
|
||
print(f"找到最新模型: {model_path}")
|
||
else:
|
||
print(f"\n错误: 找不到模型文件！")
|
||
print(f"\n可能的原因:")
|
||
print(f" 1. 还没有运行训练，先运行: python src/train.py")
|
||
print(f" 2. 使用 --model 参数指定正确路径")
|
||
return None
|
||
else:
|
||
print(f"模型路径: {model_path}")
|
||
|
||
try:
|
||
# 加载模型
|
||
model = YOLO(model_path)
|
||
|
||
# 导出为 ONNX
|
||
print(f"\n开始导出 ONNX 模型...")
|
||
print(f"输入尺寸: {imgsz}x{imgsz}")
|
||
print(f"简化模型: {simplify}")
|
||
print(f"ONNX opset: {opset}")
|
||
|
||
# 导出
|
||
export_result = model.export(
|
||
format='onnx',
|
||
imgsz=imgsz,
|
||
simplify=simplify,
|
||
opset=opset
|
||
)
|
||
|
||
if export_result:
|
||
# 重命名为 yyyy_MM_dd_HHmmss.onnx 格式
|
||
export_dir = os.path.dirname(export_result)
|
||
timestamp = datetime.datetime.now().strftime("%Y_%m_%d_%H%M%S")
|
||
new_filename = f"{timestamp}.onnx"
|
||
new_path = os.path.join(export_dir, new_filename)
|
||
|
||
# 如果文件已存在，先删除
|
||
if os.path.exists(new_path):
|
||
os.remove(new_path)
|
||
|
||
# 重命名文件
|
||
os.rename(export_result, new_path)
|
||
|
||
print(f"\n✅ 导出成功！")
|
||
print(f"ONNX 模型路径: {new_path}")
|
||
return new_path
|
||
else:
|
||
print(f"\n❌ 导出失败")
|
||
return None
|
||
|
||
except Exception as e:
|
||
print(f"\n导出失败: {e}")
|
||
return None
|
||
|
||
|
||
if __name__ == '__main__':
|
||
parser = argparse.ArgumentParser(description='将 YOLO 模型导出为 ONNX 格式')
|
||
parser.add_argument('--model', default=None,
|
||
help='模型路径（默认：自动查找最新模型）')
|
||
parser.add_argument('--imgsz', type=int, default=640,
|
||
help='输入图片尺寸（默认: 640）')
|
||
parser.add_argument('--simplify', action='store_true', default=True,
|
||
help='简化模型（默认: True）')
|
||
parser.add_argument('--no-simplify', action='store_false', dest='simplify',
|
||
help='不简化模型')
|
||
parser.add_argument('--opset', type=int, default=12,
|
||
help='ONNX opset 版本（默认: 12）')
|
||
|
||
args = parser.parse_args()
|
||
|
||
export_to_onnx(
|
||
model_path=args.model,
|
||
imgsz=args.imgsz,
|
||
simplify=args.simplify,
|
||
opset=args.opset
|
||
)
|
||
|