feat: initial HSAP platform
Huaxu Sentinel Active Safety Platform with embedded algorithm code, Docker Compose setup, and vendored dataset scaffolds for clone-and-run. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
89
algorithms/dms_yolo/code/rknn_export/convert.py
Normal file
89
algorithms/dms_yolo/code/rknn_export/convert.py
Normal file
@@ -0,0 +1,89 @@
|
||||
import sys
|
||||
import argparse
|
||||
import os
|
||||
from rknn.api import RKNN
|
||||
|
||||
|
||||
def default_dataset_path():
|
||||
# Prefer model zoo COCO subset (images included); fallback to local list.
|
||||
here = os.path.dirname(os.path.abspath(__file__))
|
||||
zoo_txt = os.path.abspath(os.path.join(here, '../../../../BK2/rknn_model_zoo-main/datasets/COCO/coco_subset_20.txt'))
|
||||
local_txt = os.path.join(here, 'datasets/COCO/coco_subset_20.txt')
|
||||
return zoo_txt if os.path.isfile(zoo_txt) else local_txt
|
||||
|
||||
|
||||
def parse_arg():
|
||||
parser = argparse.ArgumentParser(description='Convert ONNX model to RKNN format')
|
||||
|
||||
parser.add_argument('--model-path', type=str, required=True,
|
||||
help='Path to the ONNX model file')
|
||||
parser.add_argument('--platform', type=str, required=True,
|
||||
choices=['rk3562', 'rk3566', 'rk3568', 'rk3576', 'rk3588', 'rv1126b', 'rv1109', 'rv1126', 'rk1808'],
|
||||
help='Target platform. Choose from: [rk3562, rk3566, rk3568, rk3576, rk3588, rv1126b, rv1109, rv1126, rk1808]')
|
||||
parser.add_argument('--dtype', type=str, default='i8',
|
||||
choices=['i8', 'u8', 'fp'],
|
||||
help='Data type for quantization (i8/u8 for quantized, fp for float). '
|
||||
'Choose [i8, fp] for [rk3562, rk3566, rk3568, rk3576, rk3588, rv1126b]; '
|
||||
'Choose [u8, fp] for [rv1109, rv1126, rk1808]. Default: i8')
|
||||
parser.add_argument('--rknn-path', type=str, default=None,
|
||||
help='Output path for RKNN model. Default: ./<model_name>.rknn')
|
||||
parser.add_argument('--data-path', type=str, default=default_dataset_path(),
|
||||
help='Path to dataset file for quantization')
|
||||
parser.add_argument('--batch-size', type=int, default=1,
|
||||
help='Batch size for RKNN model. Default: 1')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
model_path = args.model_path
|
||||
platform = args.platform
|
||||
dataset_path = args.data_path
|
||||
batch_size = args.batch_size
|
||||
# Determine quantization based on dtype
|
||||
do_quant = args.dtype in ['i8', 'u8']
|
||||
|
||||
# Determine output path
|
||||
if args.rknn_path:
|
||||
output_path = args.rknn_path
|
||||
else:
|
||||
output_path = os.path.join('./', os.path.basename(model_path).split('.')[0] + '.rknn')
|
||||
|
||||
return model_path, platform, do_quant, output_path, dataset_path, batch_size
|
||||
|
||||
if __name__ == '__main__':
|
||||
model_path, platform, do_quant, output_path, dataset_path, batch_size = parse_arg()
|
||||
|
||||
# Create RKNN object
|
||||
rknn = RKNN(verbose=False)
|
||||
|
||||
# Pre-process config
|
||||
print('--> Config model')
|
||||
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=platform)
|
||||
print('done')
|
||||
|
||||
# Load model
|
||||
print('--> Loading model')
|
||||
ret = rknn.load_onnx(model=model_path)
|
||||
if ret != 0:
|
||||
print('Load model failed!')
|
||||
exit(ret)
|
||||
print('done')
|
||||
|
||||
# Build model
|
||||
print('--> Building model')
|
||||
ret = rknn.build(do_quantization=do_quant, dataset=dataset_path, rknn_batch_size=batch_size)
|
||||
if ret != 0:
|
||||
print('Build model failed!')
|
||||
exit(ret)
|
||||
print('done')
|
||||
|
||||
# Export rknn model
|
||||
print('--> Export rknn model')
|
||||
ret = rknn.export_rknn(output_path)
|
||||
if ret != 0:
|
||||
print('Export rknn model failed!')
|
||||
exit(ret)
|
||||
print('done')
|
||||
print('rknn model saved to: ', output_path)
|
||||
|
||||
# Release
|
||||
rknn.release()
|
||||
Reference in New Issue
Block a user