"""检查 TensorRT Engine 输出的实际 shape"""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import numpy as np

try:
    import tensorrt as trt
    import pycuda.driver as cuda
    import pycuda.autoinit
    
    engine_path = "./models/yolo11n.engine"
    
    with open(engine_path, "rb") as f:
        runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
        engine = runtime.deserialize_cuda_engine(f.read())
    
    context = engine.create_execution_context()
    
    print("=" * 60)
    print("Engine Binding Information")
    print("=" * 60)
    
    for i in range(engine.num_bindings):
        name = engine.get_binding_name(i)
        shape = engine.get_binding_shape(i)
        dtype = trt.nptype(engine.get_binding_dtype(i))
        is_input = engine.binding_is_input(i)
        
        print(f"\nBinding {i}: {name}")
        print(f"  Shape: {shape}")
        print(f"  Dtype: {dtype}")
        print(f"  Is Input: {is_input}")
    
    print("\n" + "=" * 60)
    
    input_shape = engine.get_binding_shape(0)
    output_shape = engine.get_binding_shape(1)
    
    print(f"Input shape: {input_shape}")
    print(f"Output shape: {output_shape}")
    
    input_size = np.prod([max(1, s) for s in input_shape])
    output_size = np.prod([max(1, s) for s in output_shape])
    
    print(f"Input size: {input_size}")
    print(f"Output size: {output_size}")
    
except Exception as e:
    print(f"Error: {e}")
    import traceback
    traceback.print_exc()