WIP try modeld all in python

selfdrive/modeld/constants.py

@@ -5,3 +5,6 @@ def index_function(idx, max_val=192, max_idx=32):
 
 
 T_IDXS = [index_function(idx, max_val=10.0) for idx in range(IDX_N)]
+X_IDXS = [index_function(idx, max_val=192.0) for idx in range(IDX_N)]
+lEAD_T_IDXS = [0., 2., 4., 6., 8., 10.]
+META_T_IdXS = [2., 4., 6., 8., 10.]

selfdrive/modeld/fill_model_msg.py

@@ -0,0 +1,113 @@
+import capnp
+import numpy as np
+from typing import List, Dict
+from openpilot.selfdrive.modeld.models.driving_pyx import PublishState
+from openpilot.selfdrive.modeld.constants import T_IDXS, X_IDXS, lEAD_T_IDXS, META_T_IdXS
+
+def fill_xyzt(builder, t, x, y, z, x_std=None, y_std=None, z_std=None):
+  builder.t = t
+  builder.x = x.tolist()
+  builder.y = y.tolist()
+  builder.z = z.tolist()
+  if x_std is not None: builder.xStd = x_std.tolist()
+  if y_std is not None: builder.yStd = y_std.tolist()
+  if z_std is not None: builder.zStd = z_std.tolist()
+
+def fill_xyvat(builder, t, x, y, v, a, x_std=None, y_std=None, v_std=None, a_std=None):
+  builder.t = t
+  builder.x = x.tolist()
+  builder.y = y.tolist()
+  builder.v = v.tolist()
+  builder.a = a.tolist()
+  if x_std is not None: builder.xStd = x_std.tolist()
+  if y_std is not None: builder.yStd = y_std.tolist()
+  if v_std is not None:builder.vStd = v_std.tolist()
+  if a_std is not None:builder.aStd = a_std.tolist()
+
+def fill_model_msg(msg: capnp._DynamicStructBuilder, net_output_data: Dict[str, np.ndarray], ps: PublishState,
+                   vipc_frame_id: int, vipc_frame_id_extra: int, frame_id: int, frame_drop: float,
+                   timestamp_eof: int, timestamp_llk: int, model_execution_time: float,
+                   nav_enabled: bool, valid: bool) -> None:
+  frame_age = frame_id - vipc_frame_id if frame_id > vipc_frame_id else 0
+  msg.valid = valid
+
+  modelV2 = msg.modelV2
+  modelV2.frameId = vipc_frame_id
+  modelV2.frameIdExtra = vipc_frame_id_extra
+  modelV2.frameAge = frame_age
+  modelV2.frameDropPerc = frame_drop * 100
+  modelV2.timestampEof = timestamp_eof
+  modelV2.locationMonoTime = timestamp_llk
+  modelV2.modelExecutionTime = model_execution_time
+  modelV2.navEnabled = nav_enabled
+
+  # plan
+  fill_xyzt(modelV2.position, T_IDXS, net_output_data['plan'][0,0], net_output_data['plan'][0,1], net_output_data['plan'][0,2],
+    net_output_data['plan_stds'][0,0], net_output_data['plan_stds'][0,1], net_output_data['plan_stds'][0,2])
+  fill_xyzt(modelV2.velocity, T_IDXS, net_output_data['plan'][0,3], net_output_data['plan'][0,4], net_output_data['plan'][0,5])
+  fill_xyzt(modelV2.acceleration, T_IDXS, net_output_data['plan'][0,6], net_output_data['plan'][0,7], net_output_data['plan'][0,9])
+  fill_xyzt(modelV2.orientation, T_IDXS, net_output_data['plan'][0,10], net_output_data['plan'][0,11], net_output_data['plan'][0,12])
+  fill_xyzt(modelV2.orientationRate, T_IDXS, net_output_data['plan'][0,13], net_output_data['plan'][0,14], net_output_data['plan'][0,15])
+
+  # lane lines
+  modelV2.init('laneLines', 4)
+  for i in range(4):
+    fill_xyzt(modelV2.laneLines[i], T_IDXS, np.array(X_IDXS), net_output_data['lane_lines'][0,i,:,0], net_output_data['lane_lines'][0,i,:,1])
+  modelV2.laneLineProbs = net_output_data['lane_lines_prob'][0].tolist()
+
+  # road edges
+  modelV2.init('roadEdges', 2)
+  for i in range(2):
+    fill_xyzt(modelV2.roadEdges[i], T_IDXS, np.array(X_IDXS), net_output_data['road_edges'][0,i,:,0], net_output_data['road_edges'][0,i,:,1])
+  modelV2.roadEdgeStds = net_output_data['road_edges_stds'][0,:,0,0].tolist()
+
+  # leads
+  modelV2.init('leadsV3', 2)
+  for i in range(2):
+    fill_xyvat(modelV2.leadsV3[i], lEAD_T_IDXS, net_output_data['lead'][0,i,:,0], net_output_data['lead'][0,i,:,1], net_output_data['lead'][0,i,:,2], net_output_data['lead'][0,i,:,3],
+      net_output_data['lead_stds'][0,i,:,0], net_output_data['lead_stds'][0,i,:,1], net_output_data['lead_stds'][0,i,:,2], net_output_data['lead_stds'][0,i,:,3])
+
+  # leads probs
+  # TODO
+
+  # confidence
+  # TODO
+  modelV2.confidence = 0
+
+  # meta
+  # TODO
+  modelV2.meta.engagedProb = 0.
+  modelV2.meta.hardBrakePredicted = False
+  modelV2.meta.init('disengagePredictions')
+  modelV2.meta.disengagePredictions.t = META_T_IdXS
+  modelV2.meta.disengagePredictions.brakeDisengageProbs = np.zeros(5, dtype=np.float32).tolist()
+  modelV2.meta.disengagePredictions.gasDisengageProbs = np.zeros(5, dtype=np.float32).tolist()
+  modelV2.meta.disengagePredictions.steerOverrideProbs = np.zeros(5, dtype=np.float32).tolist()
+  modelV2.meta.disengagePredictions.brake3MetersPerSecondSquaredProbs = np.zeros(5, dtype=np.float32).tolist()
+  modelV2.meta.disengagePredictions.brake4MetersPerSecondSquaredProbs = np.zeros(5, dtype=np.float32).tolist()
+  modelV2.meta.disengagePredictions.brake5MetersPerSecondSquaredProbs = np.zeros(5, dtype=np.float32).tolist()
+  modelV2.meta.desirePrediction = np.zeros(4*8, dtype=np.float32).tolist()
+
+  # temporal pose
+  modelV2.temporalPose.trans = net_output_data['sim_pose'][0,:3].tolist()
+  modelV2.temporalPose.transStd = net_output_data['sim_pose_stds'][0,:3].tolist()
+  modelV2.temporalPose.rot = net_output_data['sim_pose'][0,3:].tolist()
+  modelV2.temporalPose.rotStd = net_output_data['sim_pose_stds'][0,3:].tolist()
+
+
+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)
+  cameraOdometry = msg.cameraOdometry
+
+  cameraOdometry.frameId = vipc_frame_id
+  cameraOdometry.timestampEof = timestamp_eof
+
+  cameraOdometry.trans = net_output_data['pose'][0,:3].tolist()
+  cameraOdometry.rot = net_output_data['pose'][0,3:].tolist()
+  cameraOdometry.wideFromDeviceEuler = net_output_data['wide_from_device_euler'][0,:].tolist()
+  cameraOdometry.roadTransformTrans = net_output_data['road_transform'][0,:3].tolist()
+  cameraOdometry.transStd = net_output_data['pose_stds'][0,:3].tolist()
+  cameraOdometry.rotStd = net_output_data['pose_stds'][0,3:].tolist()
+  cameraOdometry.wideFromDeviceEulerStd = net_output_data['wide_from_device_euler_stds'][0,:].tolist()
+  cameraOdometry.roadTransformTransStd = net_output_data['road_transform_stds'][0,:3].tolist()

selfdrive/modeld/modeld.py

@@ -1,7 +1,9 @@
 #!/usr/bin/env python3
 import sys
 import time
+import pickle
 import numpy as np
+import cereal.messaging as messaging
 from pathlib import Path
 from typing import Dict, Optional
 from setproctitle import setproctitle
@@ -13,11 +15,12 @@ from openpilot.common.filter_simple import FirstOrderFilter
 from openpilot.common.realtime import config_realtime_process
 from openpilot.common.transformations.model import get_warp_matrix
 from openpilot.selfdrive.modeld.runners import ModelRunner, Runtime
+from openpilot.selfdrive.modeld.parse_model_outputs import parse_outputs
+from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_pose_msg
 from openpilot.selfdrive.modeld.models.commonmodel_pyx import ModelFrame, CLContext
 from openpilot.selfdrive.modeld.models.driving_pyx import (
-  PublishState, create_model_msg, create_pose_msg,
+  PublishState, FEATURE_LEN, HISTORY_BUFFER_LEN, DESIRE_LEN, TRAFFIC_CONVENTION_LEN, NAV_FEATURE_LEN, NAV_INSTRUCTION_LEN,
-  FEATURE_LEN, HISTORY_BUFFER_LEN, DESIRE_LEN, TRAFFIC_CONVENTION_LEN, NAV_FEATURE_LEN, NAV_INSTRUCTION_LEN,
+  NET_OUTPUT_SIZE, MODEL_FREQ)
-  OUTPUT_SIZE, NET_OUTPUT_SIZE, MODEL_FREQ)
 
 MODEL_PATHS = {
   ModelRunner.THNEED: Path(__file__).parent / 'models/supercombo.thneed',
@@ -59,8 +62,14 @@ class ModelState:
     for k,v in self.inputs.items():
       self.model.addInput(k, v)
 
+    with open(Path(__file__).parent / 'models/output_slices.pkl', 'rb') as f:
+      self.output_slices = pickle.load(f)
+
+  def slice(self, model_outputs: np.ndarray) -> Dict[str, np.ndarray]:
+    return {k: model_outputs[np.newaxis, v] for k,v in self.output_slices.items()}
+
   def run(self, buf: VisionBuf, wbuf: VisionBuf, transform: np.ndarray, transform_wide: np.ndarray,
-                inputs: Dict[str, np.ndarray], prepare_only: bool) -> Optional[np.ndarray]:
+                inputs: Dict[str, np.ndarray], prepare_only: bool) -> Optional[Dict[str, np.ndarray]]:
     # Model decides when action is completed, so desire input is just a pulse triggered on rising edge
     inputs['desire'][0] = 0
     self.inputs['desire'][:-DESIRE_LEN] = self.inputs['desire'][DESIRE_LEN:]
@@ -81,9 +90,11 @@ class ModelState:
       return None
 
     self.model.execute()
+    outputs = parse_outputs(self.slice(self.output))
+
     self.inputs['features_buffer'][:-FEATURE_LEN] = self.inputs['features_buffer'][FEATURE_LEN:]
-    self.inputs['features_buffer'][-FEATURE_LEN:] = self.output[OUTPUT_SIZE:OUTPUT_SIZE+FEATURE_LEN]
+    self.inputs['features_buffer'][-FEATURE_LEN:] = outputs['hidden_state'][0, :]
-    return self.output
+    return outputs
 
 
 def main():
@@ -130,7 +141,6 @@ def main():
   frame_id = 0
   last_vipc_frame_id = 0
   run_count = 0
-  # last = 0.0
 
   model_transform_main = np.zeros((3, 3), dtype=np.float32)
   model_transform_extra = np.zeros((3, 3), dtype=np.float32)
@@ -244,13 +254,20 @@ def main():
     model_execution_time = mt2 - mt1
 
     if model_output is not None:
-      pm.send("modelV2", create_model_msg(model_output, state, meta_main.frame_id, meta_extra.frame_id, frame_id, frame_drop_ratio,
+      modelv2_send = messaging.new_message('modelV2')
-                                          meta_main.timestamp_eof, timestamp_llk, model_execution_time, nav_enabled, live_calib_seen))
+      posenet_send = messaging.new_message('cameraOdometry')
-      pm.send("cameraOdometry", create_pose_msg(model_output, meta_main.frame_id, vipc_dropped_frames, meta_main.timestamp_eof, live_calib_seen))
+      fill_model_msg(modelv2_send, model_output, state, meta_main.frame_id, meta_extra.frame_id, frame_id, frame_drop_ratio,
+                      meta_main.timestamp_eof, timestamp_llk, model_execution_time, nav_enabled, live_calib_seen)
+
+      fill_pose_msg(posenet_send, model_output, meta_main.frame_id, vipc_dropped_frames, meta_main.timestamp_eof, live_calib_seen)
+      pm.send('modelV2', modelv2_send)
+      pm.send('cameraOdometry', posenet_send)
+      mt3 = time.perf_counter()
+      messaging_time = mt3 - mt2
+
+      print("model_execution_time: %.2fms, messaging_time: %.2fms, vipc_frame_id %u, frame_id, %u, frame_drop %.3f" %
+       ((model_execution_time)*1000, (messaging_time)*1000, meta_extra.frame_id, frame_id, frame_drop_ratio))
 
-    # print("model process: %.2fms, from last %.2fms, vipc_frame_id %u, frame_id, %u, frame_drop %.3f" %
-    #   ((mt2 - mt1)*1000, (mt1 - last)*1000, meta_extra.frame_id, frame_id, frame_drop_ratio))
-    # last = mt1
     last_vipc_frame_id = meta_main.frame_id
 
 

selfdrive/modeld/parse_model_outputs.py

@@ -0,0 +1,113 @@
+import numpy as np
+from typing import Dict
+from openpilot.selfdrive.modeld.constants import IDX_N
+
+POSE_WIDTH = 6
+SIM_POSE_WIDTH = 6
+LEAD_WIDTH = 4
+LANE_LINES_WIDTH = 2
+ROAD_EDGES_WIDTH = 2
+PLAN_WIDTH = 15
+DESIRE_PRED_WIDTH = 4
+
+NUM_LANE_LINES = 4
+NUM_ROAD_EDGES = 2
+
+LEAD_TRAJ_LEN = 6
+
+PLAN_MHP_N = 5
+LEAD_MHP_N = 2
+PLAN_MHP_SELECTION = 1
+LEAD_MHP_SELECTION = 3
+
+FCW_THRESHOLD_5MS2_HIGH = 0.15
+FCW_THRESHOLD_5MS2_LOW = 0.05
+FCW_THRESHOLD_3MS2 = 0.7
+
+def sigmoid(x):
+  return 1. / (1. + np.exp(-x))
+
+def softmax(x, axis=-1):
+  x -= np.max(x, axis=axis, keepdims=True)
+  if x.dtype == np.float32 or x.dtype == np.float64:
+    np.exp(x, out=x)
+  else:
+    x = np.exp(x)
+  x /= np.sum(x, axis=axis, keepdims=True)
+  return x
+
+def parse_mdn(name, outs, in_N=0, out_N=1, out_shape=None):
+  if name not in outs:
+    return
+  raw = outs[name]
+  raw = raw.reshape((raw.shape[0], max(in_N, 1), -1))
+
+  pred_mu = raw[:,:,:(raw.shape[2] - out_N)//2]
+  n_values = (raw.shape[2] - out_N)//2
+  pred_mu = raw[:,:,:n_values]
+  pred_std = np.exp(raw[:,:,n_values: 2*n_values])
+
+  if in_N > 1:
+    weights = np.zeros((raw.shape[0], in_N, out_N), dtype=raw.dtype)
+    for i in range(out_N):
+      weights[:,:,i - out_N] = softmax(raw[:,:,i - out_N], axis=-1)
+
+    if out_N == 1:
+      for fidx in range(weights.shape[0]):
+        idxs = np.argsort(weights[fidx][:,0])[::-1]
+        weights[fidx] = weights[fidx][idxs]
+        pred_mu[fidx] = pred_mu[fidx][idxs]
+        pred_std[fidx] = pred_std[fidx][idxs]
+    full_shape = tuple([raw.shape[0], in_N] + list(out_shape))
+    outs[name + '_weights'] = weights
+    outs[name + '_hypotheses'] = pred_mu.reshape(full_shape)
+    outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape)
+
+    pred_mu_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype)
+    pred_std_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype)
+    for fidx in range(weights.shape[0]):
+      for hidx in range(out_N):
+        idxs = np.argsort(weights[fidx,:,hidx])[::-1]
+        pred_mu_final[fidx, hidx] = pred_mu[fidx, idxs[0]]
+        pred_std_final[fidx, hidx] = pred_std[fidx, idxs[0]]
+  else:
+    pred_mu_final = pred_mu
+    pred_std_final = pred_std
+
+  if out_N > 1:
+    final_shape = tuple([raw.shape[0], out_N] + list(out_shape))
+  else:
+    final_shape = tuple([raw.shape[0],] + list(out_shape))
+  outs[name] = pred_mu_final.reshape(final_shape)
+  outs[name + '_stds'] = pred_std_final.reshape(final_shape)
+  return
+
+def parse_binary_crossentropy(name, outs):
+  if name not in outs:
+    return
+  raw = outs[name]
+  outs[name] = sigmoid(raw)
+  return
+
+def parse_categorical_crossentropy(name, outs, size=1):
+  if name not in outs:
+    return
+  raw = outs[name]
+  if size > 1:
+    raw = raw.reshape((raw.shape[0], size, -1))
+  outs[name] = softmax(raw, axis=-1)
+  return
+
+def parse_outputs(outs: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+  parse_mdn('plan', outs, in_N=PLAN_MHP_N, out_N=PLAN_MHP_SELECTION, out_shape=(IDX_N,PLAN_WIDTH))
+  parse_mdn('lane_lines', outs, in_N=0, out_N=0, out_shape=(NUM_LANE_LINES,IDX_N,LANE_LINES_WIDTH))
+  parse_mdn('road_edges', outs, in_N=0, out_N=0, out_shape=(NUM_ROAD_EDGES,IDX_N,LANE_LINES_WIDTH))
+  parse_mdn('pose', outs, in_N=0, out_N=0, out_shape=(POSE_WIDTH,))
+  parse_mdn('road_transform', outs, in_N=0, out_N=0, out_shape=(POSE_WIDTH,))
+  parse_mdn('sim_pose', outs, in_N=0, out_N=0, out_shape=(POSE_WIDTH,))
+  parse_mdn('wide_from_device_euler', outs, in_N=0, out_N=0, out_shape=(POSE_WIDTH // 2,))
+  parse_mdn('lead', outs, in_N=LEAD_MHP_N, out_N=LEAD_MHP_SELECTION, out_shape=(LEAD_TRAJ_LEN,LEAD_WIDTH))
+  for k in ['lead_prob', 'lane_lines_prob', 'meta']:
+    parse_binary_crossentropy(k, outs)
+  parse_categorical_crossentropy('desire_pred', outs, size=DESIRE_PRED_WIDTH)
+  return outs