# 加载D-Robotics 依赖库
from horizon_tc_ui import HB_ONNXRuntime
import numpy as np

# 准备模型运行的feed_dict
def prepare_input_dict(input_names):
  feed_dict = dict()
  for input_name in input_names:
      # your_custom_data_prepare代表您的自定义数据
      # 根据输入节点的类型和layout要求准备数据即可
      feed_dict[input_name] = np.random.rand(1, 1, 1, 310).astype(np.float32)
  return feed_dict

if __name__ == '__main__':
  # 创建推理Session
  sess = HB_ONNXRuntime(model_file="../Jan14_15-40-13_gym_ee_policy.onnx")
  sess_original_float = HB_ONNXRuntime(model_file='rdkx5_original_float_model.onnx')
  sess_optimized_float_model = HB_ONNXRuntime(model_file="rdkx5_optimized_float_model.onnx")
  sess_calibrated = HB_ONNXRuntime(model_file="rdkx5_calibrated_model.onnx")
  sess_quantized = HB_ONNXRuntime(model_file="rdkx5_quantized_model.onnx")

  input_names = sess.input_names
  output_names = sess.output_names
  # 准备模型输入数据
  feed_dict = prepare_input_dict(input_names)
  outputs = sess.run(output_names, feed_dict, input_offset=0)
  print("Model outputs:", outputs)
  # 开始模型推理，推理的返回值是一个list，依次与output_names指定名称一一对应
  outputs = sess_original_float.run(output_names, feed_dict, input_offset=0)
  print("Original float model outputs:", outputs)
  outputs = sess_optimized_float_model.run(output_names, feed_dict, input_offset=0)
  print("Optimized float model outputs:", outputs)
  outputs = sess_calibrated.run(output_names, feed_dict, input_offset=0)
  print("Calibrated model output:", outputs)
  outputs = sess_quantized.run(output_names, feed_dict, input_offset=0)
  print("Quantized model output:", outputs)
  
  
  
