from rknn.api import RKNN

# ONNX model uses NCHW format with float32 input normalized to 0-1
# mean_values=0, std_values=1 means the RKNN will NOT apply any normalization
# during inference — the user must provide pre-normalized (0-1) float32 NHWC input
#
# IMPORTANT: RKNN inference always expects NHWC input (1,H,W,C) regardless of
# the ONNX model's NCHW layout. The toolkit handles the transpose internally.

rknn = RKNN(verbose=True)

rknn.config(
    target_platform='rk3588',
    mean_values=[[0, 0, 0]],
    std_values=[[1, 1, 1]],
    quantized_dtype='w8a8',
    optimization_level=3
)

print('Loading ONNX model...')
ret = rknn.load_onnx(model='yolo11m_safety.onnx', input_size_list=[[3, 640, 640]])
if ret != 0: 
    print("load_onnx failed")
    exit(1)

print('Building RKNN model (do_quantization=False, float32)...')
ret = rknn.build(dataset='dataset.txt', do_quantization=False)
if ret != 0:
    print("build failed")
    exit(1)

print('Exporting RKNN model...')
rknn.export_rknn('./yolo11m_safety.rknn')

# Optional: hybrid quantization (uncomment to enable int8 weights with float outputs)
# This can improve NPU performance while keeping output precision
# ret = rknn.build(dataset='dataset.txt', do_quantization=True, quantized_dtype='asymmetric_quantized-u8')
# if ret != 0:
#     print("build with quantization failed")
#     exit(1)
# rknn.export_rknn('./yolo11m_safety_quantized.rknn')

rknn.release()
print("Done: yolo11m_safety.rknn")