diff --git a/selfdrive/modeld/constants.py b/selfdrive/modeld/constants.py index c14eac761e..ade958cd62 100644 --- a/selfdrive/modeld/constants.py +++ b/selfdrive/modeld/constants.py @@ -9,3 +9,24 @@ X_IDXS = [index_function(idx, max_val=192.0) for idx in range(IDX_N)] LEAD_T_IDXS = [0., 2., 4., 6., 8., 10.] LEAD_T_OFFSETS = [0., 2., 4.] META_T_IDXS = [2., 4., 6., 8., 10.] + +class Plan: + POSITION = slice(0, 3) + VELOCITY = slice(3, 6) + ACCELERATION = slice(6, 9) + T_FROM_CURRENT_EULER = slice(9, 12) + ORIENTATION_RATE = slice(12, 15) + +class Meta: + ENGAGED = slice(0, 1) + # next 2, 4, 6, 8, 10 seconds + GAS_DISENGAGE = slice(1, 36, 7) + BRAKE_DISENGAGE = slice(2, 36, 7) + STEER_OVERRIDE = slice(3, 36, 7) + HARD_BRAKE_3 = slice(4, 36, 7) + HARD_BRAKE_4 = slice(5, 36, 7) + HARD_BRAKE_5 = slice(6, 36, 7) + GAS_PRESS = slice(7, 36, 7) + # next 0, 2, 4, 6, 8, 10 seconds + LEFT_BLINKER = slice(36, 48, 2) + RIGHT_BLINKER = slice(37, 48, 2) \ No newline at end of file diff --git a/selfdrive/modeld/fill_model_msg.py b/selfdrive/modeld/fill_model_msg.py index 9c9e060f8c..904d633227 100644 --- a/selfdrive/modeld/fill_model_msg.py +++ b/selfdrive/modeld/fill_model_msg.py @@ -2,7 +2,9 @@ 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, LEAD_T_OFFSETS +from openpilot.selfdrive.modeld.constants import T_IDXS, X_IDXS, LEAD_T_IDXS, META_T_IDXS, LEAD_T_OFFSETS, Meta + +# TODO: use Enum Slices when possible instead of hardcoded indices def fill_xyzt(builder, t, x, y, z, x_std=None, y_std=None, z_std=None): builder.t = t @@ -42,57 +44,69 @@ def fill_model_msg(msg: capnp._DynamicStructBuilder, net_output_data: Dict[str, 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], + position = modelV2.position + fill_xyzt(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,:,8]) - fill_xyzt(modelV2.orientation, T_IDXS, net_output_data['plan'][0,:,9], net_output_data['plan'][0,:,10], net_output_data['plan'][0,:,11]) - fill_xyzt(modelV2.orientationRate, T_IDXS, net_output_data['plan'][0,:,12], net_output_data['plan'][0,:,13], net_output_data['plan'][0,:,14]) + velocity = modelV2.velocity + fill_xyzt(velocity, T_IDXS, net_output_data['plan'][0,:,3], net_output_data['plan'][0,:,4], net_output_data['plan'][0,:,5]) + acceleration = modelV2.acceleration + fill_xyzt(acceleration, T_IDXS, net_output_data['plan'][0,:,6], net_output_data['plan'][0,:,7], net_output_data['plan'][0,:,8]) + orientation = modelV2.orientation + fill_xyzt(orientation, T_IDXS, net_output_data['plan'][0,:,9], net_output_data['plan'][0,:,10], net_output_data['plan'][0,:,11]) + orientation_rate = modelV2.orientationRate + fill_xyzt(orientation_rate, T_IDXS, net_output_data['plan'][0,:,12], net_output_data['plan'][0,:,13], net_output_data['plan'][0,:,14]) # 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]) + lane_line = modelV2.laneLines[i] + fill_xyzt(lane_line, T_IDXS, np.array(X_IDXS), net_output_data['lane_lines'][0,i,:,0], net_output_data['lane_lines'][0,i,:,1]) modelV2.laneLineStds = net_output_data['lane_lines_stds'][0,:,0,0].tolist() modelV2.laneLineProbs = net_output_data['lane_lines_prob'][0,1::2].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]) + road_edge = modelV2.roadEdges[i] + fill_xyzt(road_edge, 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', 3) for i in range(3): - 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], + lead = modelV2.leadsV3[i] + fill_xyvat(lead, 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]) - modelV2.leadsV3[i].prob = net_output_data['lead_prob'][0,i].tolist() - modelV2.leadsV3[i].probTime = LEAD_T_OFFSETS[i] + lead.prob = net_output_data['lead_prob'][0,i].tolist() + lead.probTime = LEAD_T_OFFSETS[i] # 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() + meta = modelV2.meta + meta.desireState = net_output_data['desire_state'][0,:].reshape(-1).tolist() + meta.desirePrediction = net_output_data['desire_pred'][0,:].reshape(-1).tolist() + meta.engagedProb = net_output_data['meta'][0,Meta.ENGAGED].item() + meta.init('disengagePredictions') + disengage_predictions = meta.disengagePredictions + disengage_predictions.t = META_T_IDXS + disengage_predictions.brakeDisengageProbs = net_output_data['meta'][0,Meta.BRAKE_DISENGAGE].tolist() + disengage_predictions.gasDisengageProbs = net_output_data['meta'][0,Meta.GAS_DISENGAGE].tolist() + disengage_predictions.steerOverrideProbs = net_output_data['meta'][0,Meta.STEER_OVERRIDE].tolist() + disengage_predictions.brake3MetersPerSecondSquaredProbs = net_output_data['meta'][0,Meta.HARD_BRAKE_3].tolist() + disengage_predictions.brake4MetersPerSecondSquaredProbs = net_output_data['meta'][0,Meta.HARD_BRAKE_4].tolist() + disengage_predictions.brake5MetersPerSecondSquaredProbs =net_output_data['meta'][0,Meta.HARD_BRAKE_5].tolist() + + meta.hardBrakePredicted = False # 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() + temporal_pose = modelV2.temporalPose + temporal_pose.trans = net_output_data['sim_pose'][0,:3].tolist() + temporal_pose.transStd = net_output_data['sim_pose_stds'][0,:3].tolist() + temporal_pose.rot = net_output_data['sim_pose'][0,3:].tolist() + temporal_pose.rotStd = net_output_data['sim_pose_stds'][0,3:].tolist() def fill_pose_msg(msg: capnp._DynamicStructBuilder, net_output_data: Dict[str, np.ndarray], diff --git a/selfdrive/modeld/parse_model_outputs.py b/selfdrive/modeld/parse_model_outputs.py index a01122c55f..e9cc5d9d64 100755 --- a/selfdrive/modeld/parse_model_outputs.py +++ b/selfdrive/modeld/parse_model_outputs.py @@ -8,7 +8,7 @@ LEAD_WIDTH = 4 LANE_LINES_WIDTH = 2 ROAD_EDGES_WIDTH = 2 PLAN_WIDTH = 15 -DESIRE_PRED_WIDTH = 4 +DESIRE_PRED_WIDTH = 8 NUM_LANE_LINES = 4 NUM_ROAD_EDGES = 2 @@ -109,5 +109,7 @@ def parse_outputs(outs: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: 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) + for k in ['desire_pred', 'desire_state']: + parse_categorical_crossentropy(k, outs, size=DESIRE_PRED_WIDTH) + parse_categorical_crossentropy(k, outs, size=DESIRE_PRED_WIDTH) return outs