diff --git a/selfdrive/modeld/fill_model_msg.py b/selfdrive/modeld/fill_model_msg.py index dda2e41258..05a159c86d 100644 --- a/selfdrive/modeld/fill_model_msg.py +++ b/selfdrive/modeld/fill_model_msg.py @@ -1,9 +1,12 @@ +import os import capnp import numpy as np from typing import Dict from cereal import log from openpilot.selfdrive.modeld.constants import ModelConstants, Plan, Meta +SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') + ConfidenceClass = log.ModelDataV2.ConfidenceClass class PublishState: @@ -163,6 +166,10 @@ def fill_model_msg(msg: capnp._DynamicStructBuilder, net_output_data: Dict[str, else: modelV2.confidence = ConfidenceClass.red + # raw prediction if enabled + if SEND_RAW_PRED: + modelV2.rawPredictions = net_output_data['raw_pred'].tobytes() + def fill_pose_msg(msg: capnp._DynamicStructBuilder, net_output_data: Dict[str, np.ndarray], vipc_frame_id: int, vipc_dropped_frames: int, timestamp_eof: int, live_calib_seen: bool) -> None: msg.valid = live_calib_seen & (vipc_dropped_frames < 1) diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index 8bfdaaa350..5df4e5e172 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 import sys +import os import time import pickle import numpy as np @@ -20,6 +21,8 @@ from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_pose_ from openpilot.selfdrive.modeld.constants import ModelConstants from openpilot.selfdrive.modeld.models.commonmodel_pyx import ModelFrame, CLContext +SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') + MODEL_PATHS = { ModelRunner.THNEED: Path(__file__).parent / 'models/supercombo.thneed', ModelRunner.ONNX: Path(__file__).parent / 'models/supercombo.onnx'} @@ -70,7 +73,10 @@ class ModelState: self.model.addInput(k, v) def slice_outputs(self, model_outputs: np.ndarray) -> Dict[str, np.ndarray]: - return {k: model_outputs[np.newaxis, v] for k,v in self.output_slices.items()} + parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in self.output_slices.items()} + if SEND_RAW_PRED: + parsed_model_outputs['raw_pred'] = model_outputs.copy() + return parsed_model_outputs def run(self, buf: VisionBuf, wbuf: VisionBuf, transform: np.ndarray, transform_wide: np.ndarray, inputs: Dict[str, np.ndarray], prepare_only: bool) -> Optional[Dict[str, np.ndarray]]: