forked from IQ.Lvbs/IQ.Pilot
58 lines
2.4 KiB
Python
58 lines
2.4 KiB
Python
"""
|
|
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import CLMemDict, FrameDict, ModelType, NumpyDict, ShapeDict
|
|
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import ModelRunner
|
|
from openpilot.iqpilot.selfdrive.iqmodeld import MODEL_PATH
|
|
from openpilot.iqpilot.selfdrive.iqmodeld.config import ModelConstants
|
|
from openpilot.iqpilot.selfdrive.iqmodeld.parser import ArchiveParser
|
|
from openpilot.iqpilot.selfdrive.iqmodeld.runtime.ort import ORT_TYPES_TO_NP_TYPES, make_onnx_cpu_runner
|
|
|
|
|
|
def _onnx_dtype_table(session) -> dict[str, np.dtype]:
|
|
return {
|
|
tensor_info.name: ORT_TYPES_TO_NP_TYPES[tensor_info.type]
|
|
for tensor_info in session.get_inputs()
|
|
}
|
|
|
|
|
|
class ONNXRunner(ModelRunner):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.runner = make_onnx_cpu_runner(MODEL_PATH)
|
|
self._constants = ModelConstants
|
|
self._model_data = self.models.get(ModelType.supercombo)
|
|
self._input_dtypes = _onnx_dtype_table(self.runner)
|
|
self._parser = ArchiveParser()
|
|
self.parser_method_dict[ModelType.supercombo] = self._parser.parse_outputs
|
|
|
|
@property
|
|
def input_shapes(self) -> ShapeDict:
|
|
return {tensor_info.name: tensor_info.shape for tensor_info in self.runner.get_inputs()}
|
|
|
|
def _frame_as_numpy(self, stream_name: str, imgs_cl: CLMemDict, frames: FrameDict) -> np.ndarray:
|
|
flattened = frames[stream_name].as_numpy(imgs_cl[stream_name])
|
|
shaped = flattened.reshape(self.input_shapes[stream_name])
|
|
return shaped.astype(self._input_dtypes[stream_name])
|
|
|
|
def prepare_inputs(self, imgs_cl: CLMemDict, numpy_inputs: NumpyDict, frames: FrameDict) -> dict:
|
|
staged_inputs = dict(numpy_inputs)
|
|
for stream_name in imgs_cl:
|
|
staged_inputs[stream_name] = self._frame_as_numpy(stream_name, imgs_cl, frames)
|
|
self.inputs = staged_inputs
|
|
return staged_inputs
|
|
|
|
def _parse_outputs(self, model_outputs: np.ndarray) -> NumpyDict:
|
|
if self._model_data is None:
|
|
raise ValueError("Model data is not available. Ensure the model is loaded correctly.")
|
|
return self.parser_method_dict[self._model_data.model.type.raw](self._slice_outputs(model_outputs))
|
|
|
|
def _run_model(self) -> NumpyDict:
|
|
combined = self.runner.run(None, self.inputs)[0].reshape(-1)
|
|
return self._parse_outputs(combined)
|