mirror of
https://github.com/dragonpilot/dragonpilot.git
synced 2026-02-20 02:23:56 +08:00
cad17b1255
* WIP try modeld all in python * fix plan * add lane lines stds * fix lane lines prob * add lead prob * add meta * simplify plan parsing * add hard brake pred * add confidence * fix desire state and desire pred * check this file for now * rm prints * rm debug * add todos * add plan_t_idxs * same as cpp * removed cython * add wfd width - rm cpp code * add new files rm old files * get metadata at compile time * forgot this file * now uses more CPU * not used * update readme * lint * copy this too * simplify disengage probs * update model replay ref commit * update again * confidence: remove if statemens * use publish_state.enqueue * Revert "use publish_state.enqueue" This reverts commit d8807c8348338a1f773a8de00fd796abb8181404. * confidence: better shape defs * use ModelConstants class * fix confidence * Parser * slightly more power too * no inline ifs :( * confidence: just use if statements
30 lines
1.0 KiB
Python
Executable File
30 lines
1.0 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import sys
|
|
import pathlib
|
|
import onnx
|
|
import codecs
|
|
import pickle
|
|
from typing import Tuple
|
|
|
|
def get_name_and_shape(value_info:onnx.ValueInfoProto) -> Tuple[str, Tuple[int,...]]:
|
|
shape = tuple([int(dim.dim_value) for dim in value_info.type.tensor_type.shape.dim])
|
|
name = value_info.name
|
|
return name, shape
|
|
|
|
if __name__ == "__main__":
|
|
model_path = pathlib.Path(sys.argv[1])
|
|
model = onnx.load(str(model_path))
|
|
i = [x.key for x in model.metadata_props].index('output_slices')
|
|
output_slices = model.metadata_props[i].value
|
|
|
|
metadata = {}
|
|
metadata['output_slices'] = pickle.loads(codecs.decode(output_slices.encode(), "base64"))
|
|
metadata['input_shapes'] = dict([get_name_and_shape(x) for x in model.graph.input])
|
|
metadata['output_shapes'] = dict([get_name_and_shape(x) for x in model.graph.output])
|
|
|
|
metadata_path = model_path.parent / (model_path.stem + '_metadata.pkl')
|
|
with open(metadata_path, 'wb') as f:
|
|
pickle.dump(metadata, f)
|
|
|
|
print(f'saved metadata to {metadata_path}')
|