nopenpilot/selfdrive/test/process_replay/inject_model.py

84 lines
2.1 KiB
Python
Executable File

#!/usr/bin/env python3
import time
from tqdm import tqdm
import selfdrive.manager as manager
from cereal.messaging import PubMaster, recv_one, sub_sock
from tools.lib.framereader import FrameReader
def rreplace(s, old, new, occurrence):
li = s.rsplit(old, occurrence)
return new.join(li)
def regen_model(msgs, pm, frame_reader, model_sock):
# Send some livecalibration messages to initalize visiond
for msg in msgs:
if msg.which() == 'liveCalibration':
pm.send('liveCalibration', msg.as_builder())
out_msgs = []
fidx = 0
for msg in tqdm(msgs):
w = msg.which()
if w == 'frame':
msg = msg.as_builder()
img = frame_reader.get(fidx, pix_fmt="rgb24")[0][:, ::-1]
msg.frame.image = img.flatten().tobytes()
pm.send(w, msg)
model = recv_one(model_sock)
fidx += 1
out_msgs.append(model)
elif w == 'liveCalibration':
pm.send(w, msg.as_builder())
return out_msgs
def inject_model(msgs, segment_name):
if segment_name.count('--') == 2:
segment_name = rreplace(segment_name, '--', '/', 1)
frame_reader = FrameReader('cd:/'+segment_name.replace("|", "/") + "/fcamera.hevc")
manager.start_managed_process('camerad')
manager.start_managed_process('modeld')
# TODO do better than just wait for modeld to boot
time.sleep(5)
pm = PubMaster(['liveCalibration', 'frame'])
model_sock = sub_sock('model')
try:
out_msgs = regen_model(msgs, pm, frame_reader, model_sock)
except (KeyboardInterrupt, SystemExit, Exception) as e:
manager.kill_managed_process('modeld')
time.sleep(2)
manager.kill_managed_process('camerad')
raise e
manager.kill_managed_process('modeld')
time.sleep(2)
manager.kill_managed_process('camerad')
new_msgs = []
midx = 0
for msg in msgs:
if (msg.which() == 'model') and (midx < len(out_msgs)):
model = out_msgs[midx].as_builder()
model.logMonoTime = msg.logMonoTime
model = model.as_reader()
new_msgs.append(model)
midx += 1
else:
new_msgs.append(msg)
print(len(new_msgs), len(list(msgs)))
assert abs(len(new_msgs) - len(list(msgs))) < 2
return new_msgs