From b7cc6983c31b158e2552192ee6f9fdf6116adc8e Mon Sep 17 00:00:00 2001
From: YassineYousfi <yyousfi1@binghamton.edu>
Date: Fri, 20 Oct 2023 13:53:22 -0700
Subject: [PATCH] modeld: send_raw_pred under a flag (#30298)

* modeld: SEND_RAW_PRED under a flag

* not twice
---
 selfdrive/modeld/fill_model_msg.py | 7 +++++++
 selfdrive/modeld/modeld.py         | 8 +++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/selfdrive/modeld/fill_model_msg.py b/selfdrive/modeld/fill_model_msg.py
index dda2e412586240..05a159c86d77d9 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 8bfdaaa350d9c3..5df4e5e17251b6 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.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 @@ def __init__(self, context: CLContext):
       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]]: