mirror of
https://github.com/immich-app/immich.git
synced 2025-05-24 01:12:58 -04:00
54 lines
2.1 KiB
Python
54 lines
2.1 KiB
Python
import argparse
|
|
import os
|
|
|
|
parser = argparse.ArgumentParser("RKNN model converting")
|
|
parser.add_argument("model", help="Directory of the model that will be exported to RKNN ex:ViT-B-32__openai.", type=str)
|
|
parser.add_argument("target_platform", help="target platform ex:rk3566", type=str)
|
|
args = parser.parse_args()
|
|
|
|
|
|
def ConvertModel(model_path='ViT-B-32__openai/textual/model.onnx', target_platform='rk3566', dynamic_input = None):
|
|
# E build: Repeat call the 'rknn.build' or 'rknn.hybrid_quantization_step1' is not allow!
|
|
from rknn.api import RKNN
|
|
rknn = RKNN(verbose=False)
|
|
|
|
rknn.config(target_platform=target_platform, dynamic_input=dynamic_input)
|
|
ret = rknn.load_onnx(model=model_path)
|
|
|
|
if ret != 0:
|
|
print("Load failed!")
|
|
exit(ret)
|
|
|
|
ret = rknn.build(do_quantization=False)
|
|
|
|
if ret != 0:
|
|
print("Build failed!")
|
|
exit(ret)
|
|
print(model_path.replace('model.onnx',f'{target_platform}.rknn'))
|
|
ret = rknn.export_rknn(model_path.replace('model.onnx',f'{target_platform}.rknn'))
|
|
if ret != 0:
|
|
print('Export rknn model failed!')
|
|
exit(ret)
|
|
print('done')
|
|
del rknn
|
|
del RKNN
|
|
|
|
if not os.path.isfile(f'{model_path.replace("onnx","rknn")}'):
|
|
print(f'Dummy model not found at {model_path.replace("onnx","rknn")}, creating one')
|
|
with open(f'{model_path.replace("onnx","rknn")}', 'w'):
|
|
pass
|
|
|
|
|
|
if os.path.isdir(f'{args.model}/textual') and os.path.isdir(f'{args.model}/visual'): # is a clip model
|
|
print('Converting Clip model.')
|
|
ConvertModel(model_path=f'{args.model}/textual/model.onnx', target_platform=args.target_platform)
|
|
ConvertModel(model_path=f'{args.model}/visual/model.onnx', target_platform=args.target_platform)
|
|
|
|
elif os.path.isdir(f'{args.model}/detection') and os.path.isdir(f'{args.model}/recognition'): # is a facial model
|
|
print('Converting facial model.')
|
|
ConvertModel(f'{args.model}/detection/model.onnx', args.target_platform, [[[1, 3, 640, 640]]])
|
|
ConvertModel(f'{args.model}/recognition/model.onnx', args.target_platform, [[[1, 3, 112, 112]]])
|
|
|
|
else:
|
|
print('Unknown model.')
|