diff --git a/selfdrive/modeld/fill_model_msg.py b/selfdrive/modeld/fill_model_msg.py index c769668..b806bfb 100644 --- a/selfdrive/modeld/fill_model_msg.py +++ b/selfdrive/modeld/fill_model_msg.py @@ -95,10 +95,28 @@ def fill_model_msg(msg: capnp._DynamicStructBuilder, net_output_data: Dict[str, PLAN_T_IDXS[xidx] = p * ModelConstants.T_IDXS[tidx+1] + (1 - p) * ModelConstants.T_IDXS[tidx] # lane lines - modelV2.init('laneLines', 4) - for i in range(4): - lane_line = modelV2.laneLines[i] - fill_xyzt(lane_line, PLAN_T_IDXS, np.array(ModelConstants.X_IDXS), net_output_data['lane_lines'][0,i,:,0], net_output_data['lane_lines'][0,i,:,1]) + modelV2.init('laneLines', 6) + for i in range(6): + if i < 4: + lane_line = modelV2.laneLines[i] + fill_xyzt(lane_line, PLAN_T_IDXS, np.array(ModelConstants.X_IDXS), net_output_data['lane_lines'][0,i,:,0], net_output_data['lane_lines'][0,i,:,1]) + else: + lane_line = modelV2.laneLines[i] + far_lane, near_lane, road_edge = (0, 1, 0) if i == 4 else (3, 2, 1) + + y_min = net_output_data['lane_lines'][0, near_lane,:,0] + z_min = net_output_data['lane_lines'][0, near_lane,:,1] + lane_diff = np.abs(net_output_data['lane_lines'][0,near_lane] - net_output_data['lane_lines'][0,far_lane]) + road_edge_diff = np.abs(net_output_data['lane_lines'][0,near_lane] - net_output_data['road_edges'][0,road_edge]) + + y_min += np.where(lane_diff[:,0] < road_edge_diff[:,0], net_output_data['lane_lines'][0,far_lane,:,0], net_output_data['road_edges'][0,road_edge,:,0]) + z_min += np.where(lane_diff[:,1] < road_edge_diff[:,1], net_output_data['lane_lines'][0,far_lane,:,1], net_output_data['road_edges'][0,road_edge,:,1]) + + y_min /= 2 + z_min /= 2 + + fill_xyzt(lane_line, PLAN_T_IDXS, np.array(ModelConstants.X_IDXS), y_min, z_min) + modelV2.laneLineStds = net_output_data['lane_lines_stds'][0,:,0,0].tolist() modelV2.laneLineProbs = net_output_data['lane_lines_prob'][0,1::2].tolist()