#!/usr/bin/env python3
"""
导出 480x480 分辨率的 TensorRT 引擎
"""

from ultralytics import YOLO
import torch

def main():
    print("="*60)
    print("导出 480x480 TensorRT 引擎")
    print("="*60)
    
    # 加载 YOLOv11n 模型
    model_path = "C:/Users/16337/PycharmProjects/Security/yolo11n.pt"
    print(f"\n加载模型: {model_path}")
    model = YOLO(model_path)
    
    # 导出为 TensorRT 引擎
    print("\n开始导出 TensorRT 引擎...")
    print("配置:")
    print("  - 输入尺寸: 480x480")
    print("  - 精度: FP16")
    print("  - 批次大小: 动态 (1-32)")
    print()
    
    try:
        # 导出 TensorRT 引擎
        model.export(
            format='engine',
            imgsz=480,  # 480x480 分辨率
            half=True,  # FP16 精度
            dynamic=True,  # 动态批次
            batch=8,  # 优化批次大小
            workspace=4,  # 4GB workspace
            verbose=True
        )
        
        print("\n✅ TensorRT 引擎导出成功！")
        print(f"引擎文件: yolo11n.engine (480x480)")
        print("\n注意: 引擎文件会保存在当前目录")
        
    except Exception as e:
        print(f"\n❌ 导出失败: {e}")
        import traceback
        traceback.print_exc()
        return 1
    
    return 0

if __name__ == "__main__":
    import sys
    sys.exit(main())