radard: cleanup and refactor (#29071)

Re-structure + add typing

selfdrive/controls/radard.py

@@ -1,10 +1,11 @@
 #!/usr/bin/env python3
 import importlib
 import math
-from collections import defaultdict, deque
+from collections import deque
+from typing import Optional, Dict, Any
 
-import cereal.messaging as messaging
+import capnp
-from cereal import car
+from cereal import messaging, log, car
 from common.numpy_fast import interp
 from common.params import Params
 from common.realtime import Ratekeeper, Priority, config_realtime_process
@@ -17,33 +18,39 @@ from common.kalman.simple_kalman import KF1D
 _LEAD_ACCEL_TAU = 1.5
 
 # radar tracks
-SPEED, ACCEL = 0, 1   # Kalman filter states enum
+SPEED, ACCEL = 0, 1     # Kalman filter states enum
 
 # stationary qualification parameters
-v_ego_stationary = 4.   # no stationary object flag below this speed
+V_EGO_STATIONARY = 4.   # no stationary object flag below this speed
 
 RADAR_TO_CENTER = 2.7   # (deprecated) RADAR is ~ 2.7m ahead from center of car
-RADAR_TO_CAMERA = 1.52   # RADAR is ~ 1.5m ahead from center of mesh frame
+RADAR_TO_CAMERA = 1.52  # RADAR is ~ 1.5m ahead from center of mesh frame
 
 
-def get_RadarState_from_vision(lead_msg, v_ego, model_v_ego):
+class KalmanParams:
-  lead_v_rel_pred = lead_msg.v[0] - model_v_ego
+  def __init__(self, dt: float):
-  return {
+    # Lead Kalman Filter params, calculating K from A, C, Q, R requires the control library.
-    "dRel": float(lead_msg.x[0] - RADAR_TO_CAMERA),
+    # hardcoding a lookup table to compute K for values of radar_ts between 0.01s and 0.2s
-    "yRel": float(-lead_msg.y[0]),
+    assert dt > .01 and dt < .2, "Radar time step must be between .01s and 0.2s"
-    "vRel": float(lead_v_rel_pred),
+    self.A = [[1.0, dt], [0.0, 1.0]]
-    "vLead": float(v_ego + lead_v_rel_pred),
+    self.C = [1.0, 0.0]
-    "vLeadK": float(v_ego + lead_v_rel_pred),
+    #Q = np.matrix([[10., 0.0], [0.0, 100.]])
-    "aLeadK": 0.0,
+    #R = 1e3
-    "aLeadTau": 0.3,
+    #K = np.matrix([[ 0.05705578], [ 0.03073241]])
-    "fcw": False,
+    dts = [dt * 0.01 for dt in range(1, 21)]
-    "modelProb": float(lead_msg.prob),
+    K0 = [0.12287673, 0.14556536, 0.16522756, 0.18281627, 0.1988689,  0.21372394,
-    "radar": False,
+          0.22761098, 0.24069424, 0.253096,   0.26491023, 0.27621103, 0.28705801,
-    "status": True
+          0.29750003, 0.30757767, 0.31732515, 0.32677158, 0.33594201, 0.34485814,
-  }
+          0.35353899, 0.36200124]
+    K1 = [0.29666309, 0.29330885, 0.29042818, 0.28787125, 0.28555364, 0.28342219,
+          0.28144091, 0.27958406, 0.27783249, 0.27617149, 0.27458948, 0.27307714,
+          0.27162685, 0.27023228, 0.26888809, 0.26758976, 0.26633338, 0.26511557,
+          0.26393339, 0.26278425]
+    self.K = [[interp(dt, dts, K0)], [interp(dt, dts, K1)]]
 
-class Track():
+
-  def __init__(self, v_lead, kalman_params):
+class Track:
+  def __init__(self, v_lead: float, kalman_params: KalmanParams):
     self.cnt = 0
     self.aLeadTau = _LEAD_ACCEL_TAU
     self.K_A = kalman_params.A
@@ -51,7 +58,7 @@ class Track():
     self.K_K = kalman_params.K
     self.kf = KF1D([[v_lead], [0.0]], self.K_A, self.K_C, self.K_K)
 
-  def update(self, d_rel, y_rel, v_rel, v_lead, measured):
+  def update(self, d_rel: float, y_rel: float, v_rel: float, v_lead: float, measured: float):
     # relative values, copy
     self.dRel = d_rel   # LONG_DIST
     self.yRel = y_rel   # -LAT_DIST
@@ -78,13 +85,12 @@ class Track():
     # Weigh y higher since radar is inaccurate in this dimension
     return [self.dRel, self.yRel*2, self.vRel]
 
-  def reset_a_lead(self, aLeadK, aLeadTau):
+  def reset_a_lead(self, aLeadK: float, aLeadTau: float):
     self.kf = KF1D([[self.vLead], [aLeadK]], self.K_A, self.K_C, self.K_K)
     self.aLeadK = aLeadK
     self.aLeadTau = aLeadTau
 
-
+  def get_RadarState(self, model_prob: float = 0.0):
-  def get_RadarState(self, model_prob=0.0):
     return {
       "dRel": float(self.dRel),
       "yRel": float(self.yRel),
@@ -99,47 +105,25 @@ class Track():
       "aLeadTau": float(self.aLeadTau)
     }
 
-  def __str__(self):
+  def potential_low_speed_lead(self, v_ego: float):
-    ret = f"x: {self.dRel:4.1f}  y: {self.yRel:4.1f}  v: {self.vRel:4.1f}  a: {self.aLeadK:4.1f}"
-    return ret
-
-  def potential_low_speed_lead(self, v_ego):
     # stop for stuff in front of you and low speed, even without model confirmation
     # Radar points closer than 0.75, are almost always glitches on toyota radars
-    return abs(self.yRel) < 1.0 and (v_ego < v_ego_stationary) and (0.75 < self.dRel < 25)
+    return abs(self.yRel) < 1.0 and (v_ego < V_EGO_STATIONARY) and (0.75 < self.dRel < 25)
 
-  def is_potential_fcw(self, model_prob):
+  def is_potential_fcw(self, model_prob: float):
     return model_prob > .9
 
-
+  def __str__(self):
-class KalmanParams():
+    ret = f"x: {self.dRel:4.1f}  y: {self.yRel:4.1f}  v: {self.vRel:4.1f}  a: {self.aLeadK:4.1f}"
-  def __init__(self, dt):
+    return ret
-    # Lead Kalman Filter params, calculating K from A, C, Q, R requires the control library.
-    # hardcoding a lookup table to compute K for values of radar_ts between 0.01s and 0.2s
-    assert dt > .01 and dt < .2, "Radar time step must be between .01s and 0.2s"
-    self.A = [[1.0, dt], [0.0, 1.0]]
-    self.C = [1.0, 0.0]
-    #Q = np.matrix([[10., 0.0], [0.0, 100.]])
-    #R = 1e3
-    #K = np.matrix([[ 0.05705578], [ 0.03073241]])
-    dts = [dt * 0.01 for dt in range(1, 21)]
-    K0 = [0.12287673, 0.14556536, 0.16522756, 0.18281627, 0.1988689,  0.21372394,
-          0.22761098, 0.24069424, 0.253096,   0.26491023, 0.27621103, 0.28705801,
-          0.29750003, 0.30757767, 0.31732515, 0.32677158, 0.33594201, 0.34485814,
-          0.35353899, 0.36200124]
-    K1 = [0.29666309, 0.29330885, 0.29042818, 0.28787125, 0.28555364, 0.28342219,
-          0.28144091, 0.27958406, 0.27783249, 0.27617149, 0.27458948, 0.27307714,
-          0.27162685, 0.27023228, 0.26888809, 0.26758976, 0.26633338, 0.26511557,
-          0.26393339, 0.26278425]
-    self.K = [[interp(dt, dts, K0)], [interp(dt, dts, K1)]]
 
 
-def laplacian_pdf(x, mu, b):
+def laplacian_pdf(x: float, mu: float, b: float):
   b = max(b, 1e-4)
   return math.exp(-abs(x-mu)/b)
 
 
-def match_vision_to_track(v_ego, lead, tracks):
+def match_vision_to_track(v_ego: float, lead: capnp._DynamicStructReader, tracks: Dict[int, Track]):
   offset_vision_dist = lead.x[0] - RADAR_TO_CAMERA
 
   def prob(c):
@@ -162,7 +146,24 @@ def match_vision_to_track(v_ego, lead, tracks):
     return None
 
 
-def get_lead(v_ego, ready, tracks, lead_msg, model_v_ego, low_speed_override=True):
+def get_RadarState_from_vision(lead_msg: capnp._DynamicStructReader, v_ego: float, model_v_ego: float):
+  lead_v_rel_pred = lead_msg.v[0] - model_v_ego
+  return {
+    "dRel": float(lead_msg.x[0] - RADAR_TO_CAMERA),
+    "yRel": float(-lead_msg.y[0]),
+    "vRel": float(lead_v_rel_pred),
+    "vLead": float(v_ego + lead_v_rel_pred),
+    "vLeadK": float(v_ego + lead_v_rel_pred),
+    "aLeadK": 0.0,
+    "aLeadTau": 0.3,
+    "fcw": False,
+    "modelProb": float(lead_msg.prob),
+    "radar": False,
+    "status": True
+  }
+
+
+def get_lead(v_ego: float, ready: bool, tracks: Dict[int, Track], lead_msg: capnp._DynamicStructReader, model_v_ego: float, low_speed_override: bool = True) -> Dict[str, Any]:
   # Determine leads, this is where the essential logic happens
   if len(tracks) > 0 and ready and lead_msg.prob > .5:
     track = match_vision_to_track(v_ego, lead_msg, tracks)
@@ -187,22 +188,30 @@ def get_lead(v_ego, ready, tracks, lead_msg, model_v_ego, low_speed_override=Tru
   return lead_dict
 
 
-class RadarD():
+class RadarD:
-  def __init__(self, radar_ts, delay=0):
+  def __init__(self, radar_ts: float, delay: int = 0):
-    self.current_time = 0
+    self.current_time = 0.0
 
-    self.tracks = defaultdict(dict)
+    self.tracks: Dict[int, Track] = {}
     self.kalman_params = KalmanParams(radar_ts)
 
-    # v_ego
+    self.v_ego = 0.0
-    self.v_ego = 0.
+    self.v_ego_hist = deque([0.0], maxlen=delay+1)
-    self.v_ego_hist = deque([0], maxlen=delay+1)
+
+    self.radar_state: Optional[capnp._DynamicStructBuilder] = None
+    self.radar_state_valid = False
 
     self.ready = False
 
-  def update(self, sm, rr):
+  def update(self, sm: messaging.SubMaster, rr: Optional[car.RadarData]):
     self.current_time = 1e-9*max(sm.logMonoTime.values())
 
+    radar_points = []
+    radar_errors = []
+    if rr is not None:
+      radar_points = rr.points
+      radar_errors = rr.errors
+
     if sm.updated['carState']:
       self.v_ego = sm['carState'].vEgo
       self.v_ego_hist.append(self.v_ego)
@@ -210,7 +219,7 @@ class RadarD():
       self.ready = True
 
     ar_pts = {}
-    for pt in rr.points:
+    for pt in radar_points:
       ar_pts[pt.trackId] = [pt.dRel, pt.yRel, pt.vRel, pt.measured]
 
     # *** remove missing points from meta data ***
@@ -231,12 +240,11 @@ class RadarD():
       self.tracks[ids].update(rpt[0], rpt[1], rpt[2], v_lead, rpt[3])
 
     # *** publish radarState ***
-    dat = messaging.new_message('radarState')
+    self.radar_state_valid = sm.all_checks() and len(radar_errors) == 0
-    dat.valid = sm.all_checks() and len(rr.errors) == 0
+    self.radar_state = log.RadarState.new_message()
-    radarState = dat.radarState
+    self.radar_state.mdMonoTime = sm.logMonoTime['modelV2']
-    radarState.mdMonoTime = sm.logMonoTime['modelV2']
+    self.radar_state.radarErrors = list(radar_errors)
-    radarState.radarErrors = list(rr.errors)
+    self.radar_state.carStateMonoTime = sm.logMonoTime['carState']
-    radarState.carStateMonoTime = sm.logMonoTime['carState']
 
     if len(sm['modelV2'].temporalPose.trans):
       model_v_ego = sm['modelV2'].temporalPose.trans[0]
@@ -244,13 +252,32 @@ class RadarD():
       model_v_ego = self.v_ego
     leads_v3 = sm['modelV2'].leadsV3
     if len(leads_v3) > 1:
-      radarState.leadOne = get_lead(self.v_ego, self.ready, self.tracks, leads_v3[0], model_v_ego, low_speed_override=True)
+      self.radar_state.leadOne = get_lead(self.v_ego, self.ready, self.tracks, leads_v3[0], model_v_ego, low_speed_override=True)
-      radarState.leadTwo = get_lead(self.v_ego, self.ready, self.tracks, leads_v3[1], model_v_ego, low_speed_override=False)
+      self.radar_state.leadTwo = get_lead(self.v_ego, self.ready, self.tracks, leads_v3[1], model_v_ego, low_speed_override=False)
-    return dat
+
+  def publish(self, pm: messaging.PubMaster, lag_ms: float):
+    assert self.radar_state is not None
+
+    radar_msg = messaging.new_message("radarState")
+    radar_msg.valid = self.radar_state_valid
+    radar_msg.radarState = self.radar_state
+    radar_msg.radarState.cumLagMs = lag_ms
+    pm.send("radarState", radar_msg)
+
+    # publish tracks for UI debugging (keep last)
+    tracks_msg = messaging.new_message('liveTracks', len(self.tracks))
+    for index, tid in enumerate(sorted(self.tracks.keys())):
+      tracks_msg.liveTracks[index] = {
+        "trackId": tid,
+        "dRel": float(self.tracks[tid].dRel),
+        "yRel": float(self.tracks[tid].yRel),
+        "vRel": float(self.tracks[tid].vRel),
+      }
+    pm.send('liveTracks', tracks_msg)
 
 
 # fuses camera and radar data for best lead detection
-def radard_thread(sm=None, pm=None, can_sock=None):
+def radard_thread(sm: Optional[messaging.SubMaster] = None, pm: Optional[messaging.PubMaster] = None, can_sock: Optional[messaging.SubSocket] = None):
   config_realtime_process(5, Priority.CTRL_LOW)
 
   # wait for stats about the car to come in from controls
@@ -284,28 +311,13 @@ def radard_thread(sm=None, pm=None, can_sock=None):
 
     sm.update(0)
 
-    dat = RD.update(sm, rr)
+    RD.update(sm, rr)
-    dat.radarState.cumLagMs = -rk.remaining*1000.
+    RD.publish(pm, -rk.remaining*1000.0)
-
-    pm.send('radarState', dat)
-
-    # *** publish tracks for UI debugging (keep last) ***
-    tracks = RD.tracks
-    dat = messaging.new_message('liveTracks', len(tracks))
-
-    for cnt, ids in enumerate(sorted(tracks.keys())):
-      dat.liveTracks[cnt] = {
-        "trackId": ids,
-        "dRel": float(tracks[ids].dRel),
-        "yRel": float(tracks[ids].yRel),
-        "vRel": float(tracks[ids].vRel),
-      }
-    pm.send('liveTracks', dat)
 
     rk.monitor_time()
 
 
-def main(sm=None, pm=None, can_sock=None):
+def main(sm: messaging.SubMaster = None, pm: messaging.PubMaster = None, can_sock: messaging.SubSocket = None):
   radard_thread(sm, pm, can_sock)