From ee12e4b7b453c99610d8595ebbc5b9b80ec7b92d Mon Sep 17 00:00:00 2001 From: discountchubbs Date: Tue, 8 Jul 2025 06:36:23 -0700 Subject: [PATCH] Add full conditional --- .../modeld_v2/parse_model_outputs_split.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/sunnypilot/modeld_v2/parse_model_outputs_split.py b/sunnypilot/modeld_v2/parse_model_outputs_split.py index c0f1589b76..2404cf8e93 100644 --- a/sunnypilot/modeld_v2/parse_model_outputs_split.py +++ b/sunnypilot/modeld_v2/parse_model_outputs_split.py @@ -95,6 +95,9 @@ class Parser: outs[name + '_stds'] = pred_std_final.reshape(final_shape) def split_outputs(self, outs: dict[str, np.ndarray]) -> None: + if 'plan' in outs: + self.parse_mdn('plan', outs, in_N=SplitModelConstants.PLAN_MHP_N, out_N=SplitModelConstants.PLAN_MHP_SELECTION, + out_shape=(SplitModelConstants.IDX_N,SplitModelConstants.PLAN_WIDTH)) if 'lane_lines' in outs: self.parse_mdn('lane_lines', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.NUM_LANE_LINES,SplitModelConstants.IDX_N,SplitModelConstants.LANE_LINES_WIDTH)) @@ -114,25 +117,33 @@ class Parser: if 'plan' in outs: self.parse_mdn('plan', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.IDX_N,SplitModelConstants.PLAN_WIDTH)) + if 'lane_lines' in outs: + self.parse_mdn('lane_lines', outs, in_N=0, out_N=0, + out_shape=(SplitModelConstants.NUM_LANE_LINES,SplitModelConstants.IDX_N,SplitModelConstants.LANE_LINES_WIDTH)) + self.parse_mdn('road_edges', outs, in_N=0, out_N=0, + out_shape=(SplitModelConstants.NUM_ROAD_EDGES,SplitModelConstants.IDX_N,SplitModelConstants.LANE_LINES_WIDTH)) + if 'sim_pose' in outs: + self.parse_mdn('sim_pose', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,)) + for k in ['lead_prob', 'lane_lines_prob']: + self.parse_binary_crossentropy(k, outs) def parse_vision_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: self.parse_mdn('pose', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,)) self.parse_mdn('wide_from_device_euler', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.WIDE_FROM_DEVICE_WIDTH,)) self.parse_mdn('road_transform', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,)) - self.split_outputs(outs) if self.zero_inputs: self.split_zero_inputs(outs) + else: + self.split_outputs(outs) self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(SplitModelConstants.DESIRE_PRED_LEN,SplitModelConstants.DESIRE_PRED_WIDTH)) self.parse_binary_crossentropy('meta', outs) return outs def parse_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: - self.split_outputs(outs) if self.zero_inputs: self.split_zero_inputs(outs) else: - self.parse_mdn('plan', outs, in_N=SplitModelConstants.PLAN_MHP_N, out_N=SplitModelConstants.PLAN_MHP_SELECTION, - out_shape=(SplitModelConstants.IDX_N,SplitModelConstants.PLAN_WIDTH)) + self.split_outputs(outs) if 'lat_planner_solution' in outs: self.parse_mdn('lat_planner_solution', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.IDX_N,SplitModelConstants.LAT_PLANNER_SOLUTION_WIDTH)) if 'desired_curvature' in outs: