forked from IQ.Lvbs/IQ.Pilot
142 lines
4.9 KiB
Python
142 lines
4.9 KiB
Python
import numpy as np
|
|
import pyray as rl
|
|
from cereal import custom
|
|
|
|
from openpilot.selfdrive.ui.ui_state import ui_state
|
|
from openpilot.system.ui.lib.application import gui_app
|
|
from openpilot.system.ui.widgets import Widget
|
|
|
|
MODE_OFF = 0
|
|
MODE_OVERLAY = 1
|
|
MODE_REPLACE = 2
|
|
|
|
_ENV_LABEL = custom.IQEnvironment.Object.Label
|
|
_OBJECT_COLORS = {
|
|
_ENV_LABEL.car: (40, 210, 200),
|
|
_ENV_LABEL.truck: (40, 210, 200),
|
|
_ENV_LABEL.bus: (40, 210, 200),
|
|
_ENV_LABEL.motorcycle: (90, 220, 255),
|
|
_ENV_LABEL.bicycle: (90, 220, 255),
|
|
_ENV_LABEL.person: (255, 210, 90),
|
|
_ENV_LABEL.stopSign: (255, 60, 45),
|
|
_ENV_LABEL.trafficLight: (255, 190, 0),
|
|
}
|
|
|
|
_BOX_EDGES = (
|
|
(0, 1), (1, 3), (3, 2), (2, 0),
|
|
(4, 5), (5, 7), (7, 6), (6, 4),
|
|
(0, 4), (1, 5), (2, 6), (3, 7),
|
|
)
|
|
|
|
GRID_HALF_WIDTH = 12.0
|
|
GRID_MAX_DISTANCE = 90.0
|
|
GRID_STEP = 6.0
|
|
|
|
|
|
class EnvironmentRenderer(Widget):
|
|
def __init__(self):
|
|
Widget.__init__(self)
|
|
self._car_space_transform = np.zeros((3, 3), dtype=np.float32)
|
|
self._mode = MODE_OFF
|
|
self._counter = 0
|
|
|
|
def set_transform(self, transform: np.ndarray):
|
|
self._car_space_transform = transform.astype(np.float32)
|
|
|
|
def _project(self, pt: np.ndarray):
|
|
p = self._car_space_transform @ pt
|
|
if abs(p[2]) < 1e-6:
|
|
return None
|
|
return p[0] / p[2], p[1] / p[2]
|
|
|
|
def _in_rect(self, x: float, y: float) -> bool:
|
|
r = self._rect
|
|
return r.x - 400 <= x <= r.x + r.width + 400 and r.y - 400 <= y <= r.y + r.height + 400
|
|
|
|
def _render(self, rect: rl.Rectangle):
|
|
sm = ui_state.sm
|
|
if self._counter % 30 == 0:
|
|
self._mode = int(ui_state.params.get("EnvironmentView", return_default=True) or 0) if ui_state.active_bundle else 0
|
|
self._counter += 1
|
|
|
|
if self._mode == MODE_OFF:
|
|
return
|
|
if sm.recv_frame["liveCalibration"] < ui_state.started_frame:
|
|
return
|
|
|
|
if self._mode == MODE_REPLACE:
|
|
self._draw_backdrop(rect)
|
|
self._draw_ground_grid()
|
|
if sm.valid["modelV2"]:
|
|
self._draw_model_scene(sm["modelV2"])
|
|
|
|
if sm.alive["iqEnvironment"] and sm.valid["iqEnvironment"]:
|
|
self._draw_objects(sm["iqEnvironment"])
|
|
|
|
def _draw_backdrop(self, rect: rl.Rectangle):
|
|
rl.draw_rectangle_gradient_v(int(rect.x), int(rect.y), int(rect.width), int(rect.height),
|
|
rl.Color(14, 17, 22, 255), rl.Color(6, 8, 11, 255))
|
|
|
|
def _draw_ground_grid(self):
|
|
col = rl.Color(60, 70, 82, 90)
|
|
dist = GRID_STEP
|
|
while dist <= GRID_MAX_DISTANCE:
|
|
a = self._project(np.array([dist, -GRID_HALF_WIDTH, 0.0]))
|
|
b = self._project(np.array([dist, GRID_HALF_WIDTH, 0.0]))
|
|
if a and b and self._in_rect(*a) and self._in_rect(*b):
|
|
rl.draw_line_ex(rl.Vector2(*a), rl.Vector2(*b), 1.5, col)
|
|
dist += GRID_STEP
|
|
for off in np.arange(-GRID_HALF_WIDTH, GRID_HALF_WIDTH + 0.1, 3.0):
|
|
a = self._project(np.array([GRID_STEP, float(off), 0.0]))
|
|
b = self._project(np.array([GRID_MAX_DISTANCE, float(off), 0.0]))
|
|
if a and b and self._in_rect(*a) and self._in_rect(*b):
|
|
rl.draw_line_ex(rl.Vector2(*a), rl.Vector2(*b), 1.5, col)
|
|
|
|
def _draw_polyline(self, xs, ys, zs, color, thick):
|
|
pts = []
|
|
for x, y, z in zip(xs, ys, zs, strict=False):
|
|
if x < 0:
|
|
continue
|
|
s = self._project(np.array([x, y, z], dtype=np.float32))
|
|
if s and self._in_rect(*s):
|
|
pts.append(rl.Vector2(*s))
|
|
for i in range(len(pts) - 1):
|
|
rl.draw_line_ex(pts[i], pts[i + 1], thick, color)
|
|
|
|
def _draw_model_scene(self, model):
|
|
for i, lane in enumerate(model.laneLines):
|
|
a = int(np.clip(model.laneLineProbs[i], 0.0, 0.9) * 255)
|
|
self._draw_polyline(lane.x, lane.y, lane.z, rl.Color(235, 235, 235, a), 3.0)
|
|
for edge in model.roadEdges:
|
|
self._draw_polyline(edge.x, edge.y, edge.z, rl.Color(230, 70, 70, 180), 3.0)
|
|
pos = model.position
|
|
self._draw_polyline(pos.x, pos.y, pos.z, rl.Color(40, 210, 200, 220), 6.0)
|
|
|
|
def _draw_objects(self, env):
|
|
for obj in env.objects:
|
|
self._draw_box(obj)
|
|
|
|
def _draw_box(self, obj):
|
|
hx, hy = obj.length / 2.0, obj.width / 2.0
|
|
base = np.array([
|
|
[obj.x - hx, obj.y - hy, obj.z], [obj.x - hx, obj.y + hy, obj.z],
|
|
[obj.x + hx, obj.y - hy, obj.z], [obj.x + hx, obj.y + hy, obj.z],
|
|
[obj.x - hx, obj.y - hy, obj.z + obj.height], [obj.x - hx, obj.y + hy, obj.z + obj.height],
|
|
[obj.x + hx, obj.y - hy, obj.z + obj.height], [obj.x + hx, obj.y + hy, obj.z + obj.height],
|
|
], dtype=np.float32)
|
|
|
|
screen = []
|
|
for corner in base:
|
|
s = self._project(corner)
|
|
if s is None or not self._in_rect(*s):
|
|
return
|
|
screen.append(s)
|
|
|
|
r, g, b = _OBJECT_COLORS.get(obj.label, (40, 210, 200))
|
|
a = int(np.clip(obj.prob, 0.3, 1.0) * 210)
|
|
floor = [rl.Vector2(*screen[i]) for i in (0, 1, 3, 2)]
|
|
rl.draw_triangle(floor[0], floor[1], floor[2], rl.Color(r, g, b, a // 5))
|
|
rl.draw_triangle(floor[0], floor[2], floor[3], rl.Color(r, g, b, a // 5))
|
|
for i, j in _BOX_EDGES:
|
|
rl.draw_line_ex(rl.Vector2(*screen[i]), rl.Vector2(*screen[j]), 2.0, rl.Color(r, g, b, a))
|