import torch
import sys
from pathlib import Path

def export(pt_path):
    pt_path = Path(pt_path)

    # 输出 ONNX
    onnx_path = pt_path.with_suffix('.onnx')
    print(f"===> 加载 PT: {pt_path}")

    # 正确加载 YOLOv5 模型
    device = "cpu"
    ckpt = torch.load(pt_path, map_location=device)
    model = ckpt['model'].float()  # 包含完整结构
    model.eval()

    # dummy input
    dummy = torch.zeros((1, 3, 640, 640))

    print("===> 导出 ONNX...")
    torch.onnx.export(
        model,
        dummy,
        str(onnx_path),
        opset_version=12,
        input_names=['images'],
        output_names=['output'],
    )

    print(f"ONNX 导出成功: {onnx_path}")


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("用法: python 1_export_pt_to_onnx.py best.pt")
        exit(1)
    export(sys.argv[1])