mirror of
https://github.com/openmv/openmv.git
synced 2025-11-04 14:49:50 +08:00
tools/tflite2c.py: Update built-in model generator script.
Add support embedding models conditionally using the index file. Models listed in `models/index.txt` will be generated with an enable macro, which can be defined per board in `imlib_config.h` files. Other models in models/ will be built-in by default.
This commit is contained in:
parent
491fc1bd06
commit
76a17f81c9
@ -19,9 +19,15 @@ def main():
|
|||||||
parser.add_argument('--header', action = 'store_true', help = 'Generate header file.', required=False, default=False)
|
parser.add_argument('--header', action = 'store_true', help = 'Generate header file.', required=False, default=False)
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
libtf_builtin_models = []
|
tflm_builtin_models = []
|
||||||
|
tflm_builtin_models_index = []
|
||||||
print('/* NOTE: This file is auto-generated. */\n')
|
print('/* NOTE: This file is auto-generated. */\n')
|
||||||
|
|
||||||
|
with open(os.path.join(args.input, "index.txt"), 'r') as f:
|
||||||
|
for l in f.readlines():
|
||||||
|
if not l.startswith("#"):
|
||||||
|
tflm_builtin_models_index.append(os.path.basename(os.path.splitext(l.strip())[0]))
|
||||||
|
|
||||||
models_list = glob.glob(os.path.join(args.input, "*tflite"))
|
models_list = glob.glob(os.path.join(args.input, "*tflite"))
|
||||||
if (args.header):
|
if (args.header):
|
||||||
# Generate the header file
|
# Generate the header file
|
||||||
@ -32,11 +38,12 @@ def main():
|
|||||||
print(' const char **labels;')
|
print(' const char **labels;')
|
||||||
print(' const unsigned int size;')
|
print(' const unsigned int size;')
|
||||||
print(' const unsigned char *data;')
|
print(' const unsigned char *data;')
|
||||||
print('}libtf_builtin_model_t;\n')
|
print('}tflm_builtin_model_t;\n')
|
||||||
print('extern const libtf_builtin_model_t libtf_builtin_models[{:d}];'.format(len(models_list)))
|
print('extern const tflm_builtin_model_t tflm_builtin_models[];')
|
||||||
else:
|
else:
|
||||||
# Generate the C file
|
# Generate the C file
|
||||||
print('#include "libtf_builtin_models.h"')
|
print('#include "imlib_config.h"')
|
||||||
|
print('#include "tflm_builtin_models.h"')
|
||||||
for model_file in models_list:
|
for model_file in models_list:
|
||||||
model_size = os.path.getsize(model_file)
|
model_size = os.path.getsize(model_file)
|
||||||
model_name = os.path.basename(os.path.splitext(model_file)[0])
|
model_name = os.path.basename(os.path.splitext(model_file)[0])
|
||||||
@ -45,13 +52,14 @@ def main():
|
|||||||
# Generate model labels.
|
# Generate model labels.
|
||||||
labels = []
|
labels = []
|
||||||
n_labels = 0
|
n_labels = 0
|
||||||
with open(labels_file, 'r') as f:
|
if os.path.exists(labels_file):
|
||||||
labels = ['"{:s}"'.format(l.strip()) for l in f.readlines()]
|
with open(labels_file, 'r') as f:
|
||||||
n_labels = len(labels)
|
labels = ['"{:s}"'.format(l.strip()) for l in f.readlines()]
|
||||||
print('const char *libtf_{:s}_labels[] = {{{:s}}};'.format(model_name, ', '.join(labels)))
|
n_labels = len(labels)
|
||||||
|
print('static const char *tflm_{:s}_labels[] __attribute__((aligned(16))) = {{{:s}}};'.format(model_name, ', '.join(labels)))
|
||||||
|
|
||||||
# Generate model data.
|
# Generate model data.
|
||||||
print('const unsigned char libtf_{:s}_data[] = {{'.format(model_name))
|
print('static const unsigned char tflm_{:s}_data[] = {{'.format(model_name))
|
||||||
with open(model_file, 'rb') as f:
|
with open(model_file, 'rb') as f:
|
||||||
for chunk in iter(lambda: f.read(12), b''):
|
for chunk in iter(lambda: f.read(12), b''):
|
||||||
print(' ', end='')
|
print(' ', end='')
|
||||||
@ -59,18 +67,23 @@ def main():
|
|||||||
print('};')
|
print('};')
|
||||||
|
|
||||||
# Store model info in builtin models table.
|
# Store model info in builtin models table.
|
||||||
libtf_builtin_models.append([
|
tflm_builtin_models.append([
|
||||||
model_name,
|
model_name,
|
||||||
n_labels,
|
n_labels,
|
||||||
'libtf_{:s}_labels'.format(model_name),
|
'tflm_{:s}_labels'.format(model_name),
|
||||||
model_size,
|
model_size,
|
||||||
'libtf_{:s}_data'.format(model_name)]
|
'tflm_{:s}_data'.format(model_name)]
|
||||||
)
|
)
|
||||||
|
|
||||||
# Generate built-in models table.
|
# Generate built-in models table.
|
||||||
print('const libtf_builtin_model_t libtf_builtin_models[{:d}] = {{'.format(len(models_list)))
|
print('const tflm_builtin_model_t tflm_builtin_models[] = {')
|
||||||
for model in libtf_builtin_models:
|
for model in tflm_builtin_models:
|
||||||
print(' {{"{:s}", {:d}, {:s}, {:d}, {:s}}},'.format(*model))
|
if model[0] in tflm_builtin_models_index:
|
||||||
|
print(' #if defined(IMLIB_ENABLE_TFLM_BUILTIN_{:s})'.format(model[0].upper()))
|
||||||
|
print(' {{ "{:s}", {:d}, {:s}, {:d}, {:s} }},'.format(*model))
|
||||||
|
if model[0] in tflm_builtin_models_index:
|
||||||
|
print(' #endif')
|
||||||
|
print(' {0, 0, 0, 0, 0}')
|
||||||
print('};')
|
print('};')
|
||||||
|
|
||||||
if __name__ == '__main__':
|
if __name__ == '__main__':
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user