From 4069597665fd69ec399aa10f57bf15a01bfcb873 Mon Sep 17 00:00:00 2001 From: DevTekVE Date: Wed, 22 Jan 2025 08:44:22 +0100 Subject: [PATCH] Refactor curvature calculation for clarity and reuse. Introduce a dedicated `get_curvature_from_output` function to handle desired curvature retrieval, improving code readability and reusability. Replace redundant logic in curvature calculation with the new function to streamline the flow. --- selfdrive/modeld/fill_model_msg.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/selfdrive/modeld/fill_model_msg.py b/selfdrive/modeld/fill_model_msg.py index 8b09b174fd..1e4a5bd401 100644 --- a/selfdrive/modeld/fill_model_msg.py +++ b/selfdrive/modeld/fill_model_msg.py @@ -11,16 +11,25 @@ SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') ConfidenceClass = log.ModelDataV2.ConfidenceClass + def curv_from_psis(psi_target, psi_rate, vego, delay): vego = np.clip(vego, MIN_SPEED, np.inf) curv_from_psi = psi_target / (vego * delay) # epsilon to prevent divide-by-zero - return 2*curv_from_psi - psi_rate / vego + return 2 * curv_from_psi - psi_rate / vego + def get_curvature_from_plan(plan, vego, delay): psi_target = np.interp(delay, ModelConstants.T_IDXS, plan[:, Plan.T_FROM_CURRENT_EULER][:, 2]) psi_rate = plan[:, Plan.ORIENTATION_RATE][0, 2] return curv_from_psis(psi_target, psi_rate, vego, delay) + +def get_curvature_from_output(output, vego, delay): + if desired_curv := output.get('desired_curvature'): # If the model outputs the desired curvature, use that directly + return float(desired_curv[0, 0]) + + return float(get_curvature_from_plan(output['plan'][0], vego, delay)) + class PublishState: def __init__(self): self.disengage_buffer = np.zeros(ModelConstants.CONFIDENCE_BUFFER_LEN*ModelConstants.DISENGAGE_WIDTH, dtype=np.float32) @@ -77,10 +86,7 @@ def fill_model_msg(base_msg: capnp._DynamicStructBuilder, extended_msg: capnp._D extended_msg.valid = valid base_msg.valid = valid - if out_desired_curv := net_output_data.get('desired_curvature'): - desired_curv = float(out_desired_curv[0,0]) - else: - desired_curv = float(get_curvature_from_plan(net_output_data['plan'][0], v_ego, delay)) + desired_curv = float(get_curvature_from_output(net_output_data, v_ego, delay)) driving_model_data = base_msg.drivingModelData