models: simplify modeld v2 logic process (#875)

* Refactor model runner methods for improved abstraction.

Moved slicing logic to a private `_slice_outputs` method and decoupled `_run_model` for clearer subclass implementation. Removed redundant `output` attribute in `ModelState` to streamline data handling.

* Add output parsing to model_runner and remove duplicate logic

Integrates an output parser directly into `model_runner` for streamlined inference and parsing. Removes redundant parser initialization from `modeld` to avoid duplication and enhance maintainability.

* Reordering

* linter

* linter

sunnypilot/modeld_v2/model_runner.py

@@ -8,6 +8,7 @@ from openpilot.sunnypilot.modeld_v2 import MODEL_PATH, MODEL_PKL_PATH, METADATA_
 from openpilot.sunnypilot.modeld_v2.models.commonmodel_pyx import DrivingModelFrame, CLMem
 from openpilot.sunnypilot.modeld_v2.runners.ort_helpers import make_onnx_cpu_runner, ORT_TYPES_TO_NP_TYPES
 from openpilot.sunnypilot.modeld_v2.runners.tinygrad_helpers import qcom_tensor_from_opencl_address
+from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser
 from openpilot.system.hardware import TICI
 from openpilot.system.hardware.hw import Paths
 
@@ -48,6 +49,7 @@ class ModelRunner(ABC):
     self.input_shapes = self.model_metadata['input_shapes']
     self.output_slices = self.model_metadata['output_slices']
     self.inputs: dict = {}
+    self.parser = Parser()
 
   @abstractmethod
   def prepare_inputs(self, imgs_cl: dict[str, CLMem], numpy_inputs: dict[str, np.ndarray], frames: dict[str, DrivingModelFrame]) -> dict:
@@ -55,16 +57,22 @@ class ModelRunner(ABC):
     raise NotImplementedError
 
   @abstractmethod
-  def run_model(self):
+  def _run_model(self):
     """Run model inference with prepared inputs."""
+    raise NotImplementedError("This method should be implemented in subclasses.")
 
-  def slice_outputs(self, model_outputs: np.ndarray) -> dict:
+  def _slice_outputs(self, model_outputs: np.ndarray) -> dict:
     """Slice model outputs according to metadata configuration."""
     parsed_outputs = {k: model_outputs[np.newaxis, v] for k, v in self.output_slices.items()}
     if SEND_RAW_PRED:
       parsed_outputs['raw_pred'] = model_outputs.copy()
     return parsed_outputs
 
+  def run_model(self) -> dict[str, np.ndarray]:
+    """Run model inference with prepared inputs and parse outputs."""
+    result: dict[str, np.ndarray] = self.parser.parse_outputs(self._slice_outputs(self._run_model()))
+    return result
+
 
 class TinygradRunner(ModelRunner):
   """Tinygrad implementation of model runner for TICI hardware."""
@@ -108,7 +116,7 @@ class TinygradRunner(ModelRunner):
 
     return self.inputs
 
-  def run_model(self):
+  def _run_model(self):
     return self.model_run(**self.inputs).numpy().flatten()
 
 
@@ -130,5 +138,5 @@ class ONNXRunner(ModelRunner):
       self.inputs[key] = frames[key].buffer_from_cl(imgs_cl[key]).reshape(self.input_shapes[key]).astype(dtype=self.input_to_nptype[key])
     return self.inputs
 
-  def run_model(self):
+  def _run_model(self):
     return self.runner.run(None, self.inputs)[0].flatten()

sunnypilot/modeld_v2/modeld.py

@@ -18,7 +18,6 @@ from openpilot.common.transformations.camera import DEVICE_CAMERAS
 from openpilot.common.transformations.model import get_warp_matrix
 from openpilot.system import sentry
 from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper
-from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser
 from openpilot.sunnypilot.modeld_v2.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState
 from openpilot.sunnypilot.modeld_v2.constants import ModelConstants
 from openpilot.sunnypilot.modeld_v2.models.commonmodel_pyx import DrivingModelFrame, CLContext
@@ -41,7 +40,6 @@ class FrameMeta:
 class ModelState:
   frames: dict[str, DrivingModelFrame]
   inputs: dict[str, np.ndarray]
-  output: np.ndarray
   prev_desire: np.ndarray  # for tracking the rising edge of the pulse
 
   def __init__(self, context: CLContext):
@@ -64,12 +62,7 @@ class ModelState:
       if key not in self.frames: # Managed by opencl
         self.numpy_inputs[key] = np.zeros(shape, dtype=np.float32)
 
-    self.parser = Parser()
-
     if self.model_runner.is_20hz:
-      net_output_size = self.model_runner.model_metadata['output_shapes']['outputs'][1]
-      self.output = np.zeros(net_output_size, dtype=np.float32)
-
       num_elements = self.numpy_inputs['features_buffer'].shape[1]
       step_size = int(-100 / num_elements)
       self.full_features_20Hz_idxs = np.arange(step_size, step_size * (num_elements + 1), step_size)[::-1]
@@ -105,8 +98,7 @@ class ModelState:
       return None
 
     # Run model inference
-    self.output = self.model_runner.run_model()
-    outputs = self.parser.parse_outputs(self.model_runner.slice_outputs(self.output))
+    outputs = self.model_runner.run_model()
 
     if self.model_runner.is_20hz:
       self.full_features_20Hz[:-1] = self.full_features_20Hz[1:]