"""
验证 TensorRT Engine 是否可以加载
"""
import os

engine_path = r"C:\Users\16337\PycharmProjects\Security_project\yolov8n.engine"

print(f"检查文件: {engine_path}")
print(f"文件存在: {os.path.exists(engine_path)}")

if os.path.exists(engine_path):
    size = os.path.getsize(engine_path)
    print(f"文件大小: {size / (1024*1024):.2f} MB")
    
    # 检查文件头
    with open(engine_path, 'rb') as f:
        header = f.read(32)
        print(f"文件头 (hex): {header[:16].hex()}")
        # TensorRT engine 文件应该以特定的 magic number 开头
        # 不应该是 JSON 格式 (04 07 00 00 7b 22...)
    
    print("\n尝试加载 Engine...")
    try:
        import tensorrt as trt
        TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
        
        with open(engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
            engine = runtime.deserialize_cuda_engine(f.read())
        
        if engine:
            print("✓ Engine 加载成功!")
            print(f"  输入数量: {engine.num_io_tensors}")
        else:
            print("✗ Engine 加载失败 (返回 None)")
    except Exception as e:
        print(f"✗ Engine 加载失败: {e}")
else:
    print("文件不存在，请先运行:")
    print("  yolo export model=yolov8n.pt format=engine device=0 imgsz=320 batch=1 half")