Chengfang Lu
e72bc061c5
Major changes: - New frontend (platform/web/): Vite + React 18 + TypeScript + Tailwind - 4-module navigation: 数据送标 / 模型管理 / 车队管理 / 系统管理 - Data catalog with charts (DMS/ADAS/Lane 3-tab view) - Quality review workflow (标注质检): Good/Fine/Bad scoring with auto-advance - Audit enhancements: batch operations, rejection categories, Feishu notifications - Operation audit log (操作日志) - World model simulation studio (仿真工坊) - Dataset version management with snapshots and diff - ADAS 7-class dataset integration (138K images organized + compressed) - User management with Feishu integration and pagination - CRUD/search/filter on all pages, card layout redesign - PIL-optimized image overlay rendering - Auto-snapshot on build, in_review workflow stage - Removed embedded algorithm code (now in workspace)
90 lines
3.4 KiB
Python
90 lines
3.4 KiB
Python
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()
|