From 6c227a66a2a91b6c9f91d9f76c5ec268145184ca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kacper=20R=C4=85czy?= <gfw.kra@gmail.com>
Date: Mon, 30 Sep 2024 20:42:51 -0700
Subject: [PATCH] process_replay: drivingModelData migration (#33686)

* Migrate modelV2 to drivingModelData

* Add if statements

* Fix
---
 selfdrive/modeld/fill_model_msg.py         | 11 ++++++----
 selfdrive/test/process_replay/migration.py | 25 ++++++++++++++++++++++
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/selfdrive/modeld/fill_model_msg.py b/selfdrive/modeld/fill_model_msg.py
index 10a1860b58..115ec94f88 100644
--- a/selfdrive/modeld/fill_model_msg.py
+++ b/selfdrive/modeld/fill_model_msg.py
@@ -48,6 +48,12 @@ def fill_xyz_poly(builder, degree, x, y, z):
   builder.yCoefficients = coeffs[:, 1].tolist()
   builder.zCoefficients = coeffs[:, 2].tolist()
 
+def fill_lane_line_meta(builder, lane_lines, lane_line_probs):
+  builder.leftY = lane_lines[1].y[0]
+  builder.leftProb = lane_line_probs[1]
+  builder.rightY = lane_lines[2].y[0]
+  builder.rightProb = lane_line_probs[2]
+
 def fill_model_msg(base_msg: capnp._DynamicStructBuilder, extended_msg: capnp._DynamicStructBuilder,
                    net_output_data: dict[str, np.ndarray], publish_state: PublishState,
                    vipc_frame_id: int, vipc_frame_id_extra: int, frame_id: int, frame_drop: float,
@@ -130,10 +136,7 @@ def fill_model_msg(base_msg: capnp._DynamicStructBuilder, extended_msg: capnp._D
   modelV2.laneLineProbs = net_output_data['lane_lines_prob'][0,1::2].tolist()
 
   lane_line_meta = driving_model_data.laneLineMeta
-  lane_line_meta.leftY = modelV2.laneLines[1].y[0]
-  lane_line_meta.leftProb = modelV2.laneLineProbs[1]
-  lane_line_meta.rightY = modelV2.laneLines[2].y[0]
-  lane_line_meta.rightProb = modelV2.laneLineProbs[2]
+  fill_lane_line_meta(lane_line_meta, modelV2.laneLines, modelV2.laneLineProbs)
 
   # road edges
   modelV2.init('roadEdges', 2)
diff --git a/selfdrive/test/process_replay/migration.py b/selfdrive/test/process_replay/migration.py
index 9c680ada9a..c0696f2181 100644
--- a/selfdrive/test/process_replay/migration.py
+++ b/selfdrive/test/process_replay/migration.py
@@ -3,6 +3,8 @@ from collections import defaultdict
 from cereal import messaging, car
 from opendbc.car.fingerprints import MIGRATION
 from opendbc.car.toyota.values import EPS_SCALE
+from openpilot.selfdrive.modeld.constants import ModelConstants
+from openpilot.selfdrive.modeld.fill_model_msg import fill_xyz_poly, fill_lane_line_meta
 from openpilot.selfdrive.test.process_replay.vision_meta import meta_from_encode_index
 from openpilot.system.manager.process_config import managed_processes
 from panda import Panda
@@ -19,6 +21,7 @@ def migrate_all(lr, manager_states=False, panda_states=False, camera_states=Fals
   msgs = migrate_liveLocationKalman(msgs)
   msgs = migrate_liveTracks(msgs)
   msgs = migrate_driverAssistance(msgs)
+  msgs = migrate_drivingModelData(msgs)
   if manager_states:
     msgs = migrate_managerState(msgs)
   if panda_states:
@@ -29,6 +32,7 @@ def migrate_all(lr, manager_states=False, panda_states=False, camera_states=Fals
 
   return msgs
 
+
 def migrate_driverAssistance(lr):
   all_msgs = []
   for msg in lr:
@@ -39,6 +43,27 @@ def migrate_driverAssistance(lr):
       return lr
   return all_msgs
 
+
+def migrate_drivingModelData(lr):
+  all_msgs = []
+  for msg in lr:
+    all_msgs.append(msg)
+    if msg.which() == "modelV2":
+      dmd = messaging.new_message('drivingModelData', valid=msg.valid, logMonoTime=msg.logMonoTime)
+      for field in ["frameId", "frameIdExtra", "frameDropPerc", "modelExecutionTime", "action"]:
+        setattr(dmd.drivingModelData, field, getattr(msg.modelV2, field))
+      for meta_field in ["laneChangeState", "laneChangeState"]:
+        setattr(dmd.drivingModelData.meta, meta_field, getattr(msg.modelV2.meta, meta_field))
+      if len(msg.modelV2.laneLines) and len(msg.modelV2.laneLineProbs):
+        fill_lane_line_meta(dmd.drivingModelData.laneLineMeta, msg.modelV2.laneLines, msg.modelV2.laneLineProbs)
+      if all(len(a) for a in [msg.modelV2.position.x, msg.modelV2.position.y, msg.modelV2.position.z]):
+        fill_xyz_poly(dmd.drivingModelData.path, ModelConstants.POLY_PATH_DEGREE, msg.modelV2.position.x, msg.modelV2.position.y, msg.modelV2.position.z)
+      all_msgs.append(dmd.as_reader())
+    elif msg.which() == "drivingModelData":
+      return lr
+  return all_msgs
+
+
 def migrate_liveTracks(lr):
   all_msgs = []
   for msg in lr: