#!/usr/bin/env python3
"""ONNX to RKNN converter for RetinaFace"""

from rknn.api import RKNN

# Create RKNN
rknn = RKNN(verbose=False)

# Config
rknn.config(
    target_platform='rk3588',
    mean_values=[[0, 0, 0]],      # No mean subtraction
    std_values=[[1, 1, 1]],       # No normalization (already done in code)
    quantized_dtype='w8a8',       # INT8 quantization
)

# Load ONNX
print("Loading ONNX...")
rknn.load_onnx(model='RetinaFace_mobile320.onnx')

# Build
print("Building RKNN...")
rknn.build(do_quantization=True)

# Export
print("Exporting...")
rknn.export_rknn('RetinaFace_mobile320_new.rknn')
print("Done! Output: RetinaFace_mobile320_new.rknn")

rknn.release()