diff --git a/sunnypilot/modeld_v2/parse_model_outputs_split.py b/sunnypilot/modeld_v2/parse_model_outputs_split.py index 24b619ffed..4c32ba5673 100644 --- a/sunnypilot/modeld_v2/parse_model_outputs_split.py +++ b/sunnypilot/modeld_v2/parse_model_outputs_split.py @@ -118,18 +118,28 @@ class Parser: out_shape=(SplitModelConstants.IDX_N,SplitModelConstants.PLAN_WIDTH)) def split_outputs(self, outs: dict[str, np.ndarray]) -> None: + if 'desired_curvature' in outs: + self.parse_mdn('desired_curvature', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.DESIRED_CURV_WIDTH,)) + if 'desire_pred' in outs: + self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(SplitModelConstants.DESIRE_PRED_LEN,SplitModelConstants.DESIRE_PRED_WIDTH)) + if 'desire_state' in outs: + self.parse_categorical_crossentropy('desire_state', outs, out_shape=(SplitModelConstants.DESIRE_PRED_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)) + if 'lane_lines_prob' in outs: + self.parse_binary_crossentropy('lane_lines_prob', outs) + if 'lead_prob' in outs: + self.parse_binary_crossentropy('lead_prob', 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 'meta' in outs: + self.parse_binary_crossentropy('meta', outs) if 'road_edges' in outs: 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,)) - if 'lane_lines_prob' in outs: - self.parse_binary_crossentropy('lane_lines_prob', outs) - if 'lead_prob' in outs: - self.parse_binary_crossentropy('lead_prob', 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,)) @@ -137,18 +147,11 @@ class Parser: self.parse_mdn('road_transform', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,)) self.parse_dynamic_outputs(outs) 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.parse_dynamic_outputs(outs) 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: - self.parse_mdn('desired_curvature', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.DESIRED_CURV_WIDTH,)) - self.parse_categorical_crossentropy('desire_state', outs, out_shape=(SplitModelConstants.DESIRE_PRED_WIDTH,)) return outs def parse_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: