diff --git a/tools/replay/lib/ui_helpers.py b/tools/replay/lib/ui_helpers.py index 4c432a8501..9bb640679c 100644 --- a/tools/replay/lib/ui_helpers.py +++ b/tools/replay/lib/ui_helpers.py @@ -1,4 +1,4 @@ -from collections import namedtuple +import itertools from typing import Any, Dict, Tuple import matplotlib @@ -8,9 +8,8 @@ import pygame # pylint: disable=import-error from common.transformations.camera import (eon_f_frame_size, eon_f_focal_length, tici_f_frame_size, tici_f_focal_length) -from selfdrive.config import RADAR_TO_CAMERA from selfdrive.config import UIParams as UP -from tools.lib.lazy_property import lazy_property + RED = (255, 0, 0) GREEN = (0, 255, 0) @@ -19,8 +18,6 @@ YELLOW = (255, 255, 0) BLACK = (0, 0, 0) WHITE = (255, 255, 255) -_PATH_X = np.arange(192.) -_PATH_XD = np.arange(192.) _FULL_FRAME_SIZE = { } @@ -46,7 +43,24 @@ for width, height, focal in cams: METER_WIDTH = 20 -ModelUIData = namedtuple("ModelUIData", ["cpath", "lpath", "rpath", "lead", "lead_future"]) +class Calibration: + def __init__(self, num_px, extrinsic, intrinsic): + self.extrinsic = extrinsic + self.intrinsic = intrinsic + self.zoom = _BB_TO_FULL_FRAME[num_px][0, 0] + + def car_space_to_ff(self, x, y, z): + ones = np.ones_like(x) + car_space_projective = np.column_stack((x, y, z, ones)).T + + ep = self.extrinsic.dot(car_space_projective) + kep = self.intrinsic.dot(ep) + return (kep[:-1, :] / kep[-1, :]).T + + def car_space_to_bb(self, x, y, z): + pts = self.car_space_to_ff(x, y, z) + return pts / self.zoom + _COLOR_CACHE : Dict[Tuple[int, int, int], Any] = {} def find_color(lidar_surface, color): @@ -55,7 +69,6 @@ def find_color(lidar_surface, color): tcolor = 0 ret = 255 for x in lidar_surface.get_palette(): - #print tcolor, x if x[0:3] == color: ret = tcolor break @@ -63,12 +76,6 @@ def find_color(lidar_surface, color): _COLOR_CACHE[color] = ret return ret -def warp_points(pt_s, warp_matrix): - # pt_s are the source points, nxm array. - pt_d = np.dot(warp_matrix[:, :-1], pt_s.T) + warp_matrix[:, -1, None] - - # Divide by last dimension for representation in image space. - return (pt_d[:-1, :] / pt_d[-1, :]).T def to_lid_pt(y, x): px, py = -x * UP.lidar_zoom + UP.lidar_car_x, -y * UP.lidar_zoom + UP.lidar_car_y @@ -77,17 +84,10 @@ def to_lid_pt(y, x): return -1, -1 -def draw_path(y, x, color, img, calibration, top_down, lid_color=None): - # TODO: Remove big box. - uv_model_real = warp_points(np.column_stack((x, y)), calibration.car_to_model) - uv_model = np.round(uv_model_real).astype(int) - - uv_model_dots = uv_model[np.logical_and.reduce((np.all( # pylint: disable=no-member - uv_model > 0, axis=1), uv_model[:, 0] < img.shape[1] - 1, uv_model[:, 1] < - img.shape[0] - 1))] - - for i, j in ((-1, 0), (0, -1), (0, 0), (0, 1), (1, 0)): - img[uv_model_dots[:, 1] + i, uv_model_dots[:, 0] + j] = color +def draw_path(path, color, img, calibration, top_down, lid_color=None, z_off=0): + x, y, z = np.asarray(path.x), np.asarray(path.y), np.asarray(path.z) + pts = calibration.car_space_to_bb(x, -y, -z + z_off) + pts = np.round(pts).astype(int) # draw lidar path point on lidar # find color in 8 bit @@ -98,28 +98,11 @@ def draw_path(y, x, color, img, calibration, top_down, lid_color=None): if px != -1: top_down[1][px, py] = tcolor -def draw_steer_path(speed_ms, curvature, color, img, - calibration, top_down, VM, lid_color=None): - path_x = np.arange(101.) - path_y = np.multiply(path_x, np.tan(np.arcsin(np.clip(path_x * curvature, -0.999, 0.999)) / 2.)) - - draw_path(path_y, path_x, color, img, calibration, top_down, lid_color) - -def draw_lead_car(closest, top_down): - if closest is not None: - closest_y = int(round(UP.lidar_car_y - closest * UP.lidar_zoom)) - if closest_y > 0: - top_down[1][int(round(UP.lidar_car_x - METER_WIDTH * 2)):int( - round(UP.lidar_car_x + METER_WIDTH * 2)), closest_y] = find_color( - top_down[0], (255, 0, 0)) - -def draw_lead_on(img, closest_x_m, closest_y_m, calibration, color, sz=10, img_offset=(0, 0)): - uv = warp_points(np.asarray([closest_x_m, closest_y_m]), calibration.car_to_bb)[0] - u, v = int(uv[0] + img_offset[0]), int(uv[1] + img_offset[1]) - if u > 0 and u < 640 and v > 0 and v < 480 - 5: - img[v - 5 - sz:v - 5 + sz, u] = color - img[v - 5, u - sz:u + sz] = color - return u, v + height, width = img.shape[:2] + for x, y in pts: + if 1 < x < width - 1 and 1 < y < height - 1: + for a, b in itertools.permutations([-1, 0, -1], 2): + img[y + a, x + b] = color def init_plots(arr, name_to_arr_idx, plot_xlims, plot_ylims, plot_names, plot_colors, plot_styles, bigplots=False): @@ -129,7 +112,7 @@ def init_plots(arr, name_to_arr_idx, plot_xlims, plot_ylims, plot_names, plot_co "k": (0, 0, 0), "y": (1, 1, 0), "p": (0, 1, 1), - "m": (1, 0, 1) } + "m": (1, 0, 1)} if bigplots: fig = plt.figure(figsize=(6.4, 7.0)) @@ -197,97 +180,36 @@ def init_plots(arr, name_to_arr_idx, plot_xlims, plot_ylims, plot_names, plot_co return draw_plots -def draw_mpc(liveMpc, top_down): - mpc_color = find_color(top_down[0], (0, 255, 0)) - for p in zip(liveMpc.x, liveMpc.y): - px, py = to_lid_pt(*p) - top_down[1][px, py] = mpc_color - - -class CalibrationTransformsForWarpMatrix(object): - def __init__(self, num_px, model_to_full_frame, K, E): - self._model_to_full_frame = model_to_full_frame - self._K = K - self._E = E - self.num_px = num_px - - @property - def model_to_bb(self): - return _FULL_FRAME_TO_BB[self.num_px].dot(self._model_to_full_frame) - - @lazy_property - def model_to_full_frame(self): - return self._model_to_full_frame - - @lazy_property - def car_to_model(self): - return np.linalg.inv(self._model_to_full_frame).dot(self._K).dot( - self._E[:, [0, 1, 3]]) - - @lazy_property - def car_to_bb(self): - return _BB_TO_FULL_FRAME[self.num_px].dot(self._K).dot(self._E[:, [0, 1, 3]]) - - def pygame_modules_have_loaded(): return pygame.display.get_init() and pygame.font.get_init() -def draw_var(y, x, var, color, img, calibration, top_down): - # otherwise drawing gets stupid - var = max(1e-1, min(var, 0.7)) - - varcolor = tuple(np.array(color)*0.5) - draw_path(y - var, x, varcolor, img, calibration, top_down) - draw_path(y + var, x, varcolor, img, calibration, top_down) - -class ModelPoly(object): - def __init__(self, model_path): - if len(model_path.poly) == 0: - self.valid = False - return - - self.poly = np.array(model_path.poly) - - self.prob = model_path.prob - self.std = model_path.std - self.y = np.polyval(self.poly, _PATH_XD) - self.valid = True - -def extract_model_data(md): - return ModelUIData( - cpath=ModelPoly(md.path), - lpath=ModelPoly(md.leftLane), - rpath=ModelPoly(md.rightLane), - lead=md.lead, - lead_future=md.leadFuture, - ) - -def plot_model(m, VM, v_ego, curvature, imgw, calibration, top_down, d_poly, top_down_color=216): +def plot_model(m, img, calibration, top_down): if calibration is None or top_down is None: return - for lead in [m.lead, m.lead_future]: + for lead in m.leads: if lead.prob < 0.5: continue - lead_dist_from_radar = lead.dist - RADAR_TO_CAMERA - _, py_top = to_lid_pt(lead_dist_from_radar + lead.std, lead.relY) - px, py_bottom = to_lid_pt(lead_dist_from_radar - lead.std, lead.relY) - top_down[1][int(round(px - 4)):int(round(px + 4)), py_top:py_bottom] = top_down_color + x, y, _, _ = lead.xyva + x_std, _, _, _ = lead.xyvaStd + + _, py_top = to_lid_pt(x + x_std, y) + px, py_bottom = to_lid_pt(x - x_std, y) + top_down[1][int(round(px - 4)):int(round(px + 4)), py_top:py_bottom] = find_color(top_down[0], YELLOW) - color = (0, int(255 * m.lpath.prob), 0) - for path in [m.cpath, m.lpath, m.rpath]: - if path.valid: - draw_path(path.y, _PATH_XD, color, imgw, calibration, top_down, YELLOW) - draw_var(path.y, _PATH_XD, path.std, color, imgw, calibration, top_down) + for path, prob, _ in zip(m.laneLines, m.laneLineProbs, m.laneLineStds): + color = (0, int(255 * prob), 0) + draw_path(path, color, img, calibration, top_down, YELLOW, 1.22) - if d_poly is not None: - dpath_y = np.polyval(d_poly, _PATH_X) - draw_path(dpath_y, _PATH_X, RED, imgw, calibration, top_down, RED) + for edge, std in zip(m.roadEdges, m.roadEdgeStds): + prob = max(1 - std, 0) + color = (int(255 * prob), 0, 0) + draw_path(edge, color, img, calibration, top_down, RED, 1.22) - # draw user path from curvature - draw_steer_path(v_ego, curvature, BLUE, imgw, calibration, top_down, VM, BLUE) + color = (255, 0, 0) + draw_path(m.position, color, img, calibration, top_down, RED) def maybe_update_radar_points(lt, lid_overlay): diff --git a/tools/replay/sensorium.py b/tools/replay/sensorium.py deleted file mode 100755 index e49629aec2..0000000000 --- a/tools/replay/sensorium.py +++ /dev/null @@ -1,54 +0,0 @@ -#!/usr/bin/env python3 - -# Question: Can a human drive from this data? - -import os -import cv2 # pylint: disable=import-error -import numpy as np -import cereal.messaging as messaging -from common.window import Window -if os.getenv("BIG") is not None: - from common.transformations.model import BIGMODEL_INPUT_SIZE as MEDMODEL_INPUT_SIZE - from common.transformations.model import get_camera_frame_from_bigmodel_frame as get_camera_frame_from_medmodel_frame -else: - from common.transformations.model import MEDMODEL_INPUT_SIZE - from common.transformations.model import get_camera_frame_from_medmodel_frame - -from tools.replay.lib.ui_helpers import CalibrationTransformsForWarpMatrix, _FULL_FRAME_SIZE, _INTRINSICS - -if __name__ == "__main__": - sm = messaging.SubMaster(['liveCalibration']) - frame = messaging.sub_sock('frame', conflate=True) - win = Window(MEDMODEL_INPUT_SIZE[0], MEDMODEL_INPUT_SIZE[1], double=True) - num_px = 0 - calibration = None - imgff = None - - while 1: - fpkt = messaging.recv_one(frame) - if fpkt is None or len(fpkt.frame.image) == 0: - continue - sm.update(timeout=1) - rgb_img_raw = fpkt.frame.image - num_px = len(rgb_img_raw) // 3 - - if rgb_img_raw and num_px in _FULL_FRAME_SIZE.keys(): - FULL_FRAME_SIZE = _FULL_FRAME_SIZE[num_px] - imgff = np.frombuffer(rgb_img_raw, dtype=np.uint8).reshape((FULL_FRAME_SIZE[1], FULL_FRAME_SIZE[0], 3)) - imgff = imgff[:, :, ::-1] # Convert BGR to RGB - - if sm.updated['liveCalibration'] and num_px: - intrinsic_matrix = _INTRINSICS[num_px] - img_transform = np.array(fpkt.frame.transform).reshape(3, 3) - extrinsic_matrix = np.asarray(sm['liveCalibration'].extrinsicMatrix).reshape(3, 4) - ke = intrinsic_matrix.dot(extrinsic_matrix) - warp_matrix = get_camera_frame_from_medmodel_frame(ke) - calibration = CalibrationTransformsForWarpMatrix(num_px, warp_matrix, intrinsic_matrix, extrinsic_matrix) - transform = np.dot(img_transform, calibration.model_to_full_frame) - - if calibration is not None and imgff is not None: - imgw = cv2.warpAffine(imgff, transform[:2], - (MEDMODEL_INPUT_SIZE[0], MEDMODEL_INPUT_SIZE[1]), - flags=cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC) - - win.draw(imgw) diff --git a/tools/replay/ui.py b/tools/replay/ui.py index d07bf71284..a84766650a 100755 --- a/tools/replay/ui.py +++ b/tools/replay/ui.py @@ -10,34 +10,20 @@ import numpy as np import pygame # pylint: disable=import-error from common.basedir import BASEDIR -from common.transformations.model import (MODEL_CX, MODEL_CY, MODEL_INPUT_SIZE, - get_camera_frame_from_model_frame) -from selfdrive.car.toyota.interface import CarInterface as ToyotaInterface from selfdrive.config import UIParams as UP -from selfdrive.controls.lib.vehicle_model import VehicleModel import cereal.messaging as messaging from tools.replay.lib.ui_helpers import (_BB_TO_FULL_FRAME, _FULL_FRAME_SIZE, _INTRINSICS, - BLACK, BLUE, GREEN, YELLOW, RED, - CalibrationTransformsForWarpMatrix, - draw_lead_car, draw_lead_on, draw_mpc, - extract_model_data, + BLACK, GREEN, YELLOW, RED, get_blank_lid_overlay, init_plots, maybe_update_radar_points, plot_model, pygame_modules_have_loaded, - warp_points) + Calibration) os.environ['BASEDIR'] = BASEDIR ANGLE_SCALE = 5.0 def ui_thread(addr, frame_address): - # TODO: Detect car from replay and use that to select carparams - CP = ToyotaInterface.get_params("TOYOTA PRIUS 2017") - VM = VehicleModel(CP) - - CalP = np.asarray([[0, 0], [MODEL_INPUT_SIZE[0], 0], [MODEL_INPUT_SIZE[0], MODEL_INPUT_SIZE[1]], [0, MODEL_INPUT_SIZE[1]]]) - vanishing_point = np.asarray([[MODEL_CX, MODEL_CY]]) - pygame.init() pygame.font.init() assert pygame_modules_have_loaded() @@ -65,20 +51,17 @@ def ui_thread(addr, frame_address): info_font = pygame.font.SysFont("arial", 15) camera_surface = pygame.surface.Surface((640, 480), 0, 24).convert() - cameraw_surface = pygame.surface.Surface(MODEL_INPUT_SIZE, 0, 24).convert() top_down_surface = pygame.surface.Surface((UP.lidar_x, UP.lidar_y), 0, 8) frame = messaging.sub_sock('frame', addr=addr, conflate=True) sm = messaging.SubMaster(['carState', 'longitudinalPlan', 'carControl', 'radarState', 'liveCalibration', 'controlsState', - 'liveTracks', 'model', 'liveMpc', 'liveParameters', 'lateralPlan', 'frame'], addr=addr) + 'liveTracks', 'modelV2', 'liveMpc', 'liveParameters', 'lateralPlan', 'frame'], addr=addr) - calibration = None img = np.zeros((480, 640, 3), dtype='uint8') imgff = None num_px = 0 - img_transform = np.eye(3) + calibration = None - imgw = np.zeros((160, 320, 3), dtype=np.uint8) # warped image lid_overlay_blank = get_blank_lid_overlay(UP) # plots @@ -139,20 +122,14 @@ def ui_thread(addr, frame_address): imgff = np.frombuffer(rgb_img_raw, dtype=np.uint8).reshape((FULL_FRAME_SIZE[1], FULL_FRAME_SIZE[0], 3)) imgff = imgff[:, :, ::-1] # Convert BGR to RGB - cv2.warpAffine(imgff, np.dot(img_transform, _BB_TO_FULL_FRAME[num_px])[:2], - (img.shape[1], img.shape[0]), dst=img, flags=cv2.WARP_INVERSE_MAP) + zoom_matrix = _BB_TO_FULL_FRAME[num_px] + cv2.warpAffine(imgff, zoom_matrix[:2], (img.shape[1], img.shape[0]), dst=img, flags=cv2.WARP_INVERSE_MAP) intrinsic_matrix = _INTRINSICS[num_px] else: img.fill(0) intrinsic_matrix = np.eye(3) - if calibration is not None and imgff is not None: - transform = np.dot(img_transform, calibration.model_to_full_frame) - imgw = cv2.warpAffine(imgff, transform[:2], (MODEL_INPUT_SIZE[0], MODEL_INPUT_SIZE[1]), flags=cv2.WARP_INVERSE_MAP) - else: - imgw.fill(0) - sm.update() w = sm['controlsState'].lateralControlState.which() @@ -181,31 +158,15 @@ def ui_thread(addr, frame_address): plot_arr[-1, name_to_arr_idx['accel_override']] = sm['carControl'].cruiseControl.accelOverride # ***** model **** - if len(sm['model'].path.poly) > 0: - model_data = extract_model_data(sm['model']) - plot_model(model_data, VM, sm['controlsState'].vEgo, sm['controlsState'].curvature, imgw, calibration, - top_down, np.array(sm['lateralPlan'].dPolyDEPRECATED)) - - # MPC - if sm.updated['liveMpc']: - draw_mpc(sm['liveMpc'], top_down) + if sm.rcv_frame['modelV2']: + plot_model(sm['modelV2'], img, calibration, top_down) # draw all radar points maybe_update_radar_points(sm['liveTracks'], top_down[1]) if sm.updated['liveCalibration'] and num_px: extrinsic_matrix = np.asarray(sm['liveCalibration'].extrinsicMatrix).reshape(3, 4) - ke = intrinsic_matrix.dot(extrinsic_matrix) - warp_matrix = get_camera_frame_from_model_frame(ke, camera_fl=intrinsic_matrix[0][0]) - calibration = CalibrationTransformsForWarpMatrix(num_px, warp_matrix, intrinsic_matrix, extrinsic_matrix) - - # draw red pt for lead car in the main img - for lead in [sm['radarState'].leadOne, sm['radarState'].leadTwo]: - if lead.status: - if calibration is not None: - draw_lead_on(img, lead.dRel, lead.yRel, calibration, color=(192, 0, 0)) - - draw_lead_car(lead.dRel, top_down) + calibration = Calibration(num_px, extrinsic_matrix, intrinsic_matrix) # *** blits *** pygame.surfarray.blit_array(camera_surface, img.swapaxes(0, 1)) @@ -217,20 +178,11 @@ def ui_thread(addr, frame_address): screen.blit(alert_line1, (180, 150)) screen.blit(alert_line2, (180, 190)) - if calibration is not None and img is not None: - cpw = warp_points(CalP, calibration.model_to_bb) - vanishing_pointw = warp_points(vanishing_point, calibration.model_to_bb) - pygame.draw.polygon(screen, BLUE, tuple(map(tuple, cpw)), 1) - pygame.draw.circle(screen, BLUE, list(map(int, map(round, vanishing_pointw[0]))), 2) - if hor_mode: screen.blit(draw_plots(plot_arr), (640+384, 0)) else: screen.blit(draw_plots(plot_arr), (0, 600)) - pygame.surfarray.blit_array(cameraw_surface, imgw.swapaxes(0, 1)) - screen.blit(cameraw_surface, (320, 480)) - pygame.surfarray.blit_array(*top_down) screen.blit(top_down[0], (640, 0))