IQ.Pilot Release Commit @ bec7652
This commit is contained in:
269
iqpilot/selfdrive/iqmodeld/messaging.py
Normal file
269
iqpilot/selfdrive/iqmodeld/messaging.py
Normal file
@@ -0,0 +1,269 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import capnp
|
||||
import numpy as np
|
||||
|
||||
from iqpilot.cereal import log
|
||||
from iqpilot.selfdrive.iqmodeld.models.helpers import plan_x_idxs_helper
|
||||
from iqpilot.selfdrive.iqmodeld.config import ModelConstants, Plan
|
||||
from iqpilot.selfdrive.controls.lib.drive_helpers import get_curvature_from_plan
|
||||
|
||||
SEND_RAW_PRED = os.getenv("SEND_RAW_PRED")
|
||||
ConfidenceClass = log.ModelDataV2.ConfidenceClass
|
||||
|
||||
|
||||
def pick_curvature(outputs: dict[str, np.ndarray], plan_rows: np.ndarray, vehicle_speed: float,
|
||||
action_horizon: float, synthetic_lane_logic: bool) -> float:
|
||||
direct_signal = None if synthetic_lane_logic else outputs.get("desired_curvature")
|
||||
if direct_signal is not None:
|
||||
return float(direct_signal[0, 0])
|
||||
|
||||
yaw_track = plan_rows[:, Plan.T_FROM_CURRENT_EULER][:, 2]
|
||||
yaw_rate_track = plan_rows[:, Plan.ORIENTATION_RATE][:, 2]
|
||||
return float(get_curvature_from_plan(yaw_track, yaw_rate_track, ModelConstants.T_IDXS, vehicle_speed, action_horizon))
|
||||
|
||||
|
||||
@dataclass
|
||||
class DrivePacketMemory:
|
||||
disengage_rollup: np.ndarray = field(default_factory=lambda: np.zeros(
|
||||
ModelConstants.CONFIDENCE_BUFFER_LEN * ModelConstants.DISENGAGE_WIDTH, dtype=np.float32))
|
||||
brake_watch_5: np.ndarray = field(default_factory=lambda: np.zeros(
|
||||
ModelConstants.FCW_5MS2_PROBS_WIDTH, dtype=np.float32))
|
||||
brake_watch_3: np.ndarray = field(default_factory=lambda: np.zeros(
|
||||
ModelConstants.FCW_3MS2_PROBS_WIDTH, dtype=np.float32))
|
||||
|
||||
|
||||
def _assign_xyz(builder, t_points, x_track, y_track, z_track,
|
||||
x_std=None, y_std=None, z_std=None) -> None:
|
||||
builder.t = t_points
|
||||
builder.x = x_track.tolist()
|
||||
builder.y = y_track.tolist()
|
||||
builder.z = z_track.tolist()
|
||||
if x_std is not None:
|
||||
builder.xStd = x_std.tolist()
|
||||
if y_std is not None:
|
||||
builder.yStd = y_std.tolist()
|
||||
if z_std is not None:
|
||||
builder.zStd = z_std.tolist()
|
||||
|
||||
|
||||
def _assign_xyva(builder, t_points, x_track, y_track, v_track, a_track,
|
||||
x_std=None, y_std=None, v_std=None, a_std=None) -> None:
|
||||
builder.t = t_points
|
||||
builder.x = x_track.tolist()
|
||||
builder.y = y_track.tolist()
|
||||
builder.v = v_track.tolist()
|
||||
builder.a = a_track.tolist()
|
||||
if x_std is not None:
|
||||
builder.xStd = x_std.tolist()
|
||||
if y_std is not None:
|
||||
builder.yStd = y_std.tolist()
|
||||
if v_std is not None:
|
||||
builder.vStd = v_std.tolist()
|
||||
if a_std is not None:
|
||||
builder.aStd = a_std.tolist()
|
||||
|
||||
|
||||
def fill_xyz_poly(builder, degree: int, x_track: np.ndarray, y_track: np.ndarray, z_track: np.ndarray) -> None:
|
||||
stacked = np.stack([x_track, y_track, z_track], axis=1)
|
||||
coeffs = np.polynomial.polynomial.polyfit(ModelConstants.T_IDXS, stacked, deg=degree)
|
||||
builder.xCoefficients = coeffs[:, 0].tolist()
|
||||
builder.yCoefficients = coeffs[:, 1].tolist()
|
||||
builder.zCoefficients = coeffs[:, 2].tolist()
|
||||
|
||||
|
||||
def fill_lane_line_meta(builder, lane_lines, lane_probs: list[float]) -> None:
|
||||
builder.leftY = lane_lines[1].y[0]
|
||||
builder.leftProb = lane_probs[1]
|
||||
builder.rightY = lane_lines[2].y[0]
|
||||
builder.rightProb = lane_probs[2]
|
||||
|
||||
|
||||
def _roll_brake_watch(outputs: dict[str, np.ndarray], memory: DrivePacketMemory, meta_layout) -> bool:
|
||||
memory.brake_watch_5[:-1] = memory.brake_watch_5[1:]
|
||||
memory.brake_watch_5[-1] = outputs["meta"][0, meta_layout.HARD_BRAKE_5][0]
|
||||
memory.brake_watch_3[:-1] = memory.brake_watch_3[1:]
|
||||
memory.brake_watch_3[-1] = outputs["meta"][0, meta_layout.HARD_BRAKE_3][0]
|
||||
return bool(
|
||||
(memory.brake_watch_5 > ModelConstants.FCW_THRESHOLDS_5MS2).all()
|
||||
and (memory.brake_watch_3 > ModelConstants.FCW_THRESHOLDS_3MS2).all()
|
||||
)
|
||||
|
||||
|
||||
def _confidence_bucket(outputs: dict[str, np.ndarray], memory: DrivePacketMemory, meta_layout, frame_id: int):
|
||||
width = ModelConstants.DISENGAGE_WIDTH
|
||||
if frame_id % (2 * ModelConstants.MODEL_FREQ) == 0:
|
||||
brake_probs = outputs["meta"][0, meta_layout.BRAKE_DISENGAGE]
|
||||
gas_probs = outputs["meta"][0, meta_layout.GAS_DISENGAGE]
|
||||
steer_probs = outputs["meta"][0, meta_layout.STEER_OVERRIDE]
|
||||
takeover_curve = 1 - ((1 - brake_probs) * (1 - gas_probs) * (1 - steer_probs))
|
||||
independent = np.r_[takeover_curve[0], np.diff(takeover_curve) / (1 - takeover_curve[:-1])]
|
||||
memory.disengage_rollup[:-width] = memory.disengage_rollup[width:]
|
||||
memory.disengage_rollup[-width:] = independent
|
||||
|
||||
score = 0.0
|
||||
for idx in range(width):
|
||||
score += memory.disengage_rollup[idx * width + width - 1 - idx].item() / width
|
||||
|
||||
if score < ModelConstants.RYG_GREEN:
|
||||
return ConfidenceClass.green
|
||||
if score < ModelConstants.RYG_YELLOW:
|
||||
return ConfidenceClass.yellow
|
||||
return ConfidenceClass.red
|
||||
|
||||
|
||||
def _write_plan_family(model_packet, driving_packet, outputs: dict[str, np.ndarray]) -> None:
|
||||
plan_rows = outputs["plan"][0]
|
||||
plan_stds = outputs["plan_stds"][0]
|
||||
_assign_xyz(model_packet.position, ModelConstants.T_IDXS, *plan_rows[:, Plan.POSITION].T, *plan_stds[:, Plan.POSITION].T)
|
||||
_assign_xyz(model_packet.velocity, ModelConstants.T_IDXS, *plan_rows[:, Plan.VELOCITY].T)
|
||||
_assign_xyz(model_packet.acceleration, ModelConstants.T_IDXS, *plan_rows[:, Plan.ACCELERATION].T)
|
||||
_assign_xyz(model_packet.orientation, ModelConstants.T_IDXS, *plan_rows[:, Plan.T_FROM_CURRENT_EULER].T)
|
||||
_assign_xyz(model_packet.orientationRate, ModelConstants.T_IDXS, *plan_rows[:, Plan.ORIENTATION_RATE].T)
|
||||
fill_xyz_poly(driving_packet.path, ModelConstants.POLY_PATH_DEGREE, *plan_rows[:, Plan.POSITION].T)
|
||||
|
||||
|
||||
def _write_temporal_pose(model_packet, outputs: dict[str, np.ndarray]) -> None:
|
||||
pose_packet = model_packet.temporalPoseDEPRECATED
|
||||
if "sim_pose" in outputs:
|
||||
half_width = ModelConstants.POSE_WIDTH // 2
|
||||
pose_packet.trans = outputs["sim_pose"][0, :half_width].tolist()
|
||||
pose_packet.transStd = outputs["sim_pose_stds"][0, :half_width].tolist()
|
||||
pose_packet.rot = outputs["sim_pose"][0, half_width:].tolist()
|
||||
pose_packet.rotStd = outputs["sim_pose_stds"][0, half_width:].tolist()
|
||||
return
|
||||
|
||||
pose_packet.trans = outputs["plan"][0, 0, Plan.VELOCITY].tolist()
|
||||
pose_packet.transStd = outputs["plan_stds"][0, 0, Plan.VELOCITY].tolist()
|
||||
pose_packet.rot = outputs["plan"][0, 0, Plan.ORIENTATION_RATE].tolist()
|
||||
pose_packet.rotStd = outputs["plan_stds"][0, 0, Plan.ORIENTATION_RATE].tolist()
|
||||
|
||||
|
||||
def _write_lane_family(model_packet, driving_packet, outputs: dict[str, np.ndarray]) -> None:
|
||||
time_axis = plan_x_idxs_helper(ModelConstants, Plan, outputs)
|
||||
model_packet.init("laneLines", 4)
|
||||
for lane_idx in range(4):
|
||||
lane_builder = model_packet.laneLines[lane_idx]
|
||||
_assign_xyz(
|
||||
lane_builder,
|
||||
time_axis,
|
||||
np.array(ModelConstants.X_IDXS),
|
||||
outputs["lane_lines"][0, lane_idx, :, 0],
|
||||
outputs["lane_lines"][0, lane_idx, :, 1],
|
||||
)
|
||||
model_packet.laneLineStds = outputs["lane_lines_stds"][0, :, 0, 0].tolist()
|
||||
model_packet.laneLineProbs = outputs["lane_lines_prob"][0, 1::2].tolist()
|
||||
fill_lane_line_meta(driving_packet.laneLineMeta, model_packet.laneLines, model_packet.laneLineProbs)
|
||||
|
||||
model_packet.init("roadEdges", 2)
|
||||
for edge_idx in range(2):
|
||||
edge_builder = model_packet.roadEdges[edge_idx]
|
||||
_assign_xyz(
|
||||
edge_builder,
|
||||
time_axis,
|
||||
np.array(ModelConstants.X_IDXS),
|
||||
outputs["road_edges"][0, edge_idx, :, 0],
|
||||
outputs["road_edges"][0, edge_idx, :, 1],
|
||||
)
|
||||
model_packet.roadEdgeStds = outputs["road_edges_stds"][0, :, 0, 0].tolist()
|
||||
|
||||
|
||||
def _write_leads(model_packet, outputs: dict[str, np.ndarray]) -> None:
|
||||
model_packet.init("leadsV3", 3)
|
||||
for lead_idx in range(3):
|
||||
lead_builder = model_packet.leadsV3[lead_idx]
|
||||
_assign_xyva(
|
||||
lead_builder,
|
||||
ModelConstants.LEAD_T_IDXS,
|
||||
*outputs["lead"][0, lead_idx].T,
|
||||
*outputs["lead_stds"][0, lead_idx].T,
|
||||
)
|
||||
lead_builder.prob = outputs["lead_prob"][0, lead_idx].tolist()
|
||||
lead_builder.probTime = ModelConstants.LEAD_T_OFFSETS[lead_idx]
|
||||
|
||||
|
||||
def _write_meta(model_packet, outputs: dict[str, np.ndarray], memory: DrivePacketMemory, meta_layout, frame_id: int) -> None:
|
||||
meta = model_packet.meta
|
||||
meta.desireState = outputs["desire_state"][0].reshape(-1).tolist()
|
||||
meta.desirePrediction = outputs["desire_pred"][0].reshape(-1).tolist()
|
||||
meta.engagedProb = outputs["meta"][0, meta_layout.ENGAGED].item()
|
||||
meta.init("disengagePredictions")
|
||||
|
||||
pred = meta.disengagePredictions
|
||||
pred.t = ModelConstants.META_T_IDXS
|
||||
pred.brakeDisengageProbs = outputs["meta"][0, meta_layout.BRAKE_DISENGAGE].tolist()
|
||||
pred.gasDisengageProbs = outputs["meta"][0, meta_layout.GAS_DISENGAGE].tolist()
|
||||
pred.steerOverrideProbs = outputs["meta"][0, meta_layout.STEER_OVERRIDE].tolist()
|
||||
pred.brake3MetersPerSecondSquaredProbs = outputs["meta"][0, meta_layout.HARD_BRAKE_3].tolist()
|
||||
pred.brake4MetersPerSecondSquaredProbs = outputs["meta"][0, meta_layout.HARD_BRAKE_4].tolist()
|
||||
pred.brake5MetersPerSecondSquaredProbs = outputs["meta"][0, meta_layout.HARD_BRAKE_5].tolist()
|
||||
|
||||
if hasattr(meta_layout, "GAS_PRESS") and hasattr(meta_layout, "BRAKE_PRESS"):
|
||||
pred.gasPressProbs = outputs["meta"][0, meta_layout.GAS_PRESS].tolist()
|
||||
pred.brakePressProbs = outputs["meta"][0, meta_layout.BRAKE_PRESS].tolist()
|
||||
|
||||
meta.hardBrakePredicted = _roll_brake_watch(outputs, memory, meta_layout)
|
||||
model_packet.confidence = _confidence_bucket(outputs, memory, meta_layout, frame_id)
|
||||
|
||||
|
||||
def populate_drive_messages(primary_msg: capnp._DynamicStructBuilder, extended_msg: capnp._DynamicStructBuilder,
|
||||
outputs: dict[str, np.ndarray], action: log.ModelDataV2.Action,
|
||||
memory: DrivePacketMemory, vipc_frame_id: int, vipc_frame_id_extra: int,
|
||||
frame_id: int, frame_drop: float, timestamp_eof: int,
|
||||
model_execution_time: float, valid: bool, meta_layout) -> None:
|
||||
frame_age = frame_id - vipc_frame_id if frame_id > vipc_frame_id else 0
|
||||
frame_drop_percent = frame_drop * 100
|
||||
primary_msg.valid = valid
|
||||
extended_msg.valid = valid
|
||||
|
||||
driving_packet = primary_msg.drivingModelData
|
||||
driving_packet.frameId = vipc_frame_id
|
||||
driving_packet.frameIdExtra = vipc_frame_id_extra
|
||||
driving_packet.frameDropPerc = frame_drop_percent
|
||||
driving_packet.modelExecutionTime = model_execution_time
|
||||
driving_packet.action = action
|
||||
|
||||
model_packet = extended_msg.modelV2
|
||||
model_packet.frameId = vipc_frame_id
|
||||
model_packet.frameIdExtra = vipc_frame_id_extra
|
||||
model_packet.frameAge = frame_age
|
||||
model_packet.frameDropPerc = frame_drop_percent
|
||||
model_packet.timestampEof = timestamp_eof
|
||||
model_packet.modelExecutionTime = model_execution_time
|
||||
model_packet.action = action
|
||||
|
||||
_write_plan_family(model_packet, driving_packet, outputs)
|
||||
_write_temporal_pose(model_packet, outputs)
|
||||
_write_lane_family(model_packet, driving_packet, outputs)
|
||||
_write_leads(model_packet, outputs)
|
||||
_write_meta(model_packet, outputs, memory, meta_layout, vipc_frame_id)
|
||||
|
||||
if SEND_RAW_PRED:
|
||||
model_packet.rawPredictions = outputs["raw_pred"].tobytes()
|
||||
|
||||
|
||||
def populate_odometry_message(msg: capnp._DynamicStructBuilder, outputs: dict[str, np.ndarray],
|
||||
vipc_frame_id: int, vipc_dropped_frames: int,
|
||||
timestamp_eof: int, live_calib_seen: bool) -> None:
|
||||
msg.valid = live_calib_seen & (vipc_dropped_frames < 1)
|
||||
odo = msg.cameraOdometry
|
||||
odo.frameId = vipc_frame_id
|
||||
odo.timestampEof = timestamp_eof
|
||||
odo.trans = outputs["pose"][0, :3].tolist()
|
||||
odo.rot = outputs["pose"][0, 3:].tolist()
|
||||
odo.wideFromDeviceEuler = outputs["wide_from_device_euler"][0, :].tolist()
|
||||
odo.roadTransformTrans = outputs["road_transform"][0, :3].tolist()
|
||||
odo.transStd = outputs["pose_stds"][0, :3].tolist()
|
||||
odo.rotStd = outputs["pose_stds"][0, 3:].tolist()
|
||||
odo.wideFromDeviceEulerStd = outputs["wide_from_device_euler_stds"][0, :].tolist()
|
||||
odo.roadTransformTransStd = outputs["road_transform_stds"][0, :3].tolist()
|
||||
|
||||
__all__ = [
|
||||
"DrivePacketMemory",
|
||||
"pick_curvature",
|
||||
"populate_drive_messages",
|
||||
"populate_odometry_message",
|
||||
]
|
||||
Reference in New Issue
Block a user