diff --git a/system/ui/lib/shader_polygon.py b/system/ui/lib/shader_polygon.py index ab8b4473bb..cc4cc38456 100644 --- a/system/ui/lib/shader_polygon.py +++ b/system/ui/lib/shader_polygon.py @@ -1,9 +1,13 @@ +import os import pyray as rl import numpy as np +from raylib import rl as raw_rl from dataclasses import dataclass from typing import Any, Optional, cast from openpilot.system.ui.lib.application import gui_app, GL_VERSION +USE_VERTEX_GRADIENTS = os.getenv("USE_VERTEX_GRADIENTS", "1") == "1" + MAX_GRADIENT_COLORS = 20 # includes stops as well @@ -200,6 +204,56 @@ def triangulate(pts: np.ndarray) -> np.ndarray: return tri_strip +def _draw_strip_vertex_gradient(tri_strip: np.ndarray, gradient: Gradient, origin_rect: rl.Rectangle) -> None: + n = len(tri_strip) + if n < 3 or not gradient.colors: + return + + start_x = origin_rect.x + gradient.start[0] * origin_rect.width + start_y = origin_rect.y + gradient.start[1] * origin_rect.height + end_x = origin_rect.x + gradient.end[0] * origin_rect.width + end_y = origin_rect.y + gradient.end[1] * origin_rect.height + + dx = start_x - end_x + dy = start_y - end_y + len2 = max(dx * dx + dy * dy, 1e-6) + + vx = tri_strip[:, 0] - end_x + vy = tri_strip[:, 1] - end_y + t = np.clip((vx * dx + vy * dy) / len2, 0.0, 1.0) + + colors = gradient.colors + stops = gradient.stops if len(gradient.stops) == len(colors) else [i / max(1, len(colors) - 1) for i in range(len(colors))] + r_stops = [c.r for c in colors] + g_stops = [c.g for c in colors] + b_stops = [c.b for c in colors] + a_stops = [c.a for c in colors] + + r = np.interp(t, stops, r_stops).astype(np.uint8) + g = np.interp(t, stops, g_stops).astype(np.uint8) + b = np.interp(t, stops, b_stops).astype(np.uint8) + a = np.interp(t, stops, a_stops).astype(np.uint8) + + rl_begin = raw_rl.rlBegin + rl_end = raw_rl.rlEnd + rl_color = raw_rl.rlColor4ub + rl_vertex = raw_rl.rlVertex2f + + rl_begin(4) # RL_TRIANGLES + for i in range(2, n): + if i % 2 == 0: + i0, i1, i2 = i, i - 2, i - 1 + else: + i0, i1, i2 = i, i - 1, i - 2 + rl_color(r[i0], g[i0], b[i0], a[i0]) + rl_vertex(tri_strip[i0, 0], tri_strip[i0, 1]) + rl_color(r[i1], g[i1], b[i1], a[i1]) + rl_vertex(tri_strip[i1, 0], tri_strip[i1, 1]) + rl_color(r[i2], g[i2], b[i2], a[i2]) + rl_vertex(tri_strip[i2, 0], tri_strip[i2, 1]) + rl_end() + + def draw_polygon(origin_rect: rl.Rectangle, points: np.ndarray, color: Optional[rl.Color] = None, gradient: Gradient | None = None): @@ -215,12 +269,19 @@ def draw_polygon(origin_rect: rl.Rectangle, points: np.ndarray, # Triangulate via interleaving tri_strip = triangulate(pts) - vertices = rl.ffi.from_buffer("Vector2 *", tri_strip) + if len(tri_strip) < 3: + return if gradient is None: + vertices = rl.ffi.from_buffer("Vector2 *", tri_strip) rl.draw_triangle_strip(vertices, len(tri_strip), color or rl.WHITE) return + if USE_VERTEX_GRADIENTS: + _draw_strip_vertex_gradient(tri_strip, gradient, origin_rect) + return + + vertices = rl.ffi.from_buffer("Vector2 *", tri_strip) state = ShaderState.get_instance() state.initialize() diff --git a/system/ui/lib/tests/test_shader_polygon.py b/system/ui/lib/tests/test_shader_polygon.py index faf6e88dc8..ee745c67ff 100644 --- a/system/ui/lib/tests/test_shader_polygon.py +++ b/system/ui/lib/tests/test_shader_polygon.py @@ -59,7 +59,8 @@ def test_triangulate_preserves_coordinates_in_contiguous_float_buffer(points): @pytest.mark.parametrize("use_gradient", [False, True]) @pytest.mark.parametrize("point_count", [4, 5, 66]) -def test_draw_polygon_passes_native_vertices(monkeypatch, use_gradient, point_count): +def test_draw_polygon_legacy_shader_fallback(monkeypatch, use_gradient, point_count): + monkeypatch.setattr(shader_polygon, "USE_VERTEX_GRADIENTS", False) points = np.arange(point_count * 4, dtype=np.float64).reshape(-1, 2)[::2] rect = rl.Rectangle(0, 0, 100, 100) color = rl.Color(10, 20, 30, 40) @@ -84,3 +85,76 @@ def test_draw_polygon_passes_native_vertices(monkeypatch, use_gradient, point_co assert calls == (["initialize", ("configure", (state, None, gradient, rect)), ("begin", "shader"), "draw", "end"] if use_gradient else ["draw"]) + + +@pytest.mark.parametrize("use_gradient", [False, True]) +@pytest.mark.parametrize("point_count", [4, 5, 66]) +def test_draw_polygon_native_vertex_gradients(monkeypatch, use_gradient, point_count): + monkeypatch.setattr(shader_polygon, "USE_VERTEX_GRADIENTS", True) + points = np.arange(point_count * 4, dtype=np.float64).reshape(-1, 2)[::2] + rect = rl.Rectangle(0, 0, 100, 100) + color = rl.Color(10, 20, 30, 40) + gradient = Gradient((0, 0), (1, 1), [color, rl.Color(*rl.WHITE)], [0, 1]) if use_gradient else None + + raw_calls = [] + emitted_vertices = [] + emitted_colors = [] + fake_raw = SimpleNamespace( + rlBegin=lambda m: raw_calls.append(("rlBegin", m)), + rlEnd=lambda: raw_calls.append("rlEnd"), + rlColor4ub=lambda r, g, b, a: emitted_colors.append((int(r), int(g), int(b), int(a))), + rlVertex2f=lambda x, y: emitted_vertices.append((float(x), float(y))), + ) + monkeypatch.setattr(shader_polygon, "raw_rl", fake_raw) + + draw_calls = [] + def draw_strip(vertices, count, fill): + assert rl.ffi.typeof(vertices) == rl.ffi.typeof("Vector2 *") + coordinates = [[vertices[i].x, vertices[i].y] for i in range(count)] + np.testing.assert_array_equal(coordinates, triangulate(points)) + assert fill == color + draw_calls.append("draw") + + monkeypatch.setattr(rl, "draw_triangle_strip", draw_strip) + + draw_polygon(rect, points, gradient=gradient, color=None if use_gradient else color) + + if use_gradient: + strip = triangulate(points) + n = len(strip) + assert raw_calls == [("rlBegin", 4), "rlEnd"] + expected_vertex_count = (n - 2) * 3 + assert len(emitted_vertices) == expected_vertex_count + assert len(emitted_colors) == expected_vertex_count + assert draw_calls == [] + else: + assert draw_calls == ["draw"] + assert raw_calls == [] + + +def test_vertex_gradient_color_interpolation(monkeypatch): + monkeypatch.setattr(shader_polygon, "USE_VERTEX_GRADIENTS", True) + points = np.array([ + [0.0, 100.0], + [0.0, 0.0], + [10.0, 0.0], + [10.0, 100.0], + ], dtype=np.float32) + rect = rl.Rectangle(0, 0, 10, 100) + color_start = rl.Color(255, 0, 0, 255) + color_end = rl.Color(0, 255, 0, 128) + gradient = Gradient((0.0, 1.0), (0.0, 0.0), [color_start, color_end], [0.0, 1.0]) + + emitted_colors = [] + fake_raw = SimpleNamespace( + rlBegin=lambda m: None, + rlEnd=lambda: None, + rlColor4ub=lambda r, g, b, a: emitted_colors.append((int(r), int(g), int(b), int(a))), + rlVertex2f=lambda x, y: None, + ) + monkeypatch.setattr(shader_polygon, "raw_rl", fake_raw) + + draw_polygon(rect, points, gradient=gradient) + + assert (255, 0, 0, 255) in emitted_colors + assert (0, 255, 0, 128) in emitted_colors