import tensorrt as trt
import pycuda.driver as cuda
import numpy as np

'''
    ultralytics导出的engine文件 根本用不了啊。反序列化的engine一直为 None，莫名奇妙
'''
print(trt.__version__)
# 1.构建日志记录器
logger = trt.Logger(trt.Logger.INFO)

# 2.反序列化
runtime = trt.Runtime(logger)

engine_path = '/home/admin-root/haotian/python哈汽安全帽识别/5000_HeadHelmet/train2/weights/best.engine'
# engine_path = '/home/admin-root/220_project/tensorrtx-master/yolov8/build/220version0826.engine'
with open(engine_path, 'rb') as f:
    serialized_engine = f.read()

engine = runtime.deserialize_cuda_engine(serialized_engine)

# 打印模型信息
for idx in range(engine.num_bindings):
    name = engine.get_tensor_name(idx)
    is_input = engine.get_tensor_mode(name)
    op_type = engine.get_tensor_dtype(name)
    shape = engine.get_tensor_shape(name)
    print('input id: ',idx, '\\tis input: ', is_input, '\\tbinding name: ', name, '\\tshape: ', shape, '\\ttype: ', op_type)
# 3.执行推理
# 3.1上下文 、
context = engine.create_execution_context()
# 为输入输出指定缓冲区