|
| 1 | +#!/usr/bin/env python3 |
| 2 | + |
| 3 | +# Copyright 2023 TIER IV, Inc. |
| 4 | +# |
| 5 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 6 | +# you may not use this file except in compliance with the License. |
| 7 | +# You may obtain a copy of the License at |
| 8 | +# |
| 9 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +# |
| 11 | +# Unless required by applicable law or agreed to in writing, software |
| 12 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 13 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 | +# See the License for the specific language governing permissions and |
| 15 | +# limitations under the License. |
| 16 | + |
| 17 | +import argparse |
| 18 | +from dataclasses import dataclass |
| 19 | +from itertools import cycle |
| 20 | +import math |
| 21 | +from threading import Lock |
| 22 | +import time |
| 23 | + |
| 24 | +import imageio |
| 25 | +import matplotlib |
| 26 | +import matplotlib.pyplot as plt |
| 27 | +import numpy as np |
| 28 | +import rclpy |
| 29 | +from rclpy.node import Node |
| 30 | +from tier4_debug_msgs.msg import Float64MultiArrayStamped |
| 31 | + |
| 32 | +matplotlib.use("TKAgg") |
| 33 | + |
| 34 | + |
| 35 | +@dataclass |
| 36 | +class NPC: |
| 37 | + x: float |
| 38 | + y: float |
| 39 | + th: float |
| 40 | + width: float |
| 41 | + height: float |
| 42 | + speed: float |
| 43 | + dangerous: bool |
| 44 | + ref_object_enter_time: float |
| 45 | + ref_object_exit_time: float |
| 46 | + collision_start_time: float |
| 47 | + collision_start_dist: float |
| 48 | + collision_end_time: float |
| 49 | + collision_end_dist: float |
| 50 | + pred_x: list[float] |
| 51 | + pred_y: list[float] |
| 52 | + |
| 53 | + def __init__(self, data: list[float]): |
| 54 | + self.x = data[0] |
| 55 | + self.y = data[1] |
| 56 | + self.th = data[2] |
| 57 | + self.width = data[3] |
| 58 | + self.height = data[4] |
| 59 | + self.speed = data[5] |
| 60 | + self.dangerous = bool(int(data[6])) |
| 61 | + self.ref_object_enter_time = data[7] |
| 62 | + self.ref_object_exit_time = data[8] |
| 63 | + self.collision_start_time = data[9] |
| 64 | + self.collision_start_dist = data[10] |
| 65 | + self.collision_end_time = data[11] |
| 66 | + self.collision_end_dist = data[12] |
| 67 | + self.first_collision_x = data[13] |
| 68 | + self.first_collision_y = data[14] |
| 69 | + self.last_collision_x = data[15] |
| 70 | + self.last_collision_y = data[16] |
| 71 | + self.pred_x = data[17:58:2] |
| 72 | + self.pred_y = data[18:58:2] |
| 73 | + |
| 74 | + |
| 75 | +class TTCVisualizer(Node): |
| 76 | + def __init__(self, args): |
| 77 | + super().__init__("ttc_visualizer") |
| 78 | + self.ttc_dist_pub = self.create_subscription( |
| 79 | + Float64MultiArrayStamped, |
| 80 | + "/planning/scenario_planning/lane_driving/behavior_planning/behavior_velocity_planner/debug/intersection/ego_ttc", |
| 81 | + self.on_ego_ttc, |
| 82 | + 1, |
| 83 | + ) |
| 84 | + self.ttc_time_pub = self.create_subscription( |
| 85 | + Float64MultiArrayStamped, |
| 86 | + "/planning/scenario_planning/lane_driving/behavior_planning/behavior_velocity_planner/debug/intersection/object_ttc", |
| 87 | + self.on_object_ttc, |
| 88 | + 1, |
| 89 | + ) |
| 90 | + self.args = args |
| 91 | + self.lane_id = args.lane_id |
| 92 | + self.ego_ttc_data = None |
| 93 | + self.object_ttc_data = None |
| 94 | + self.npc_vehicles = [] |
| 95 | + self.images = [] |
| 96 | + self.last_sub = time.time() |
| 97 | + |
| 98 | + self.plot_timer = self.create_timer(0.2, self.on_plot_timer) |
| 99 | + self.fig = plt.figure(figsize=(13, 6)) |
| 100 | + self.ttc_ax = self.fig.add_subplot(1, 2, 1) |
| 101 | + self.ttc_vel_ax = self.ttc_ax.twinx() |
| 102 | + self.world_ax = self.fig.add_subplot(1, 2, 2) |
| 103 | + self.lock = Lock() |
| 104 | + self.color_list = [ |
| 105 | + "#e41a1c", |
| 106 | + "#377eb8", |
| 107 | + "#4daf4a", |
| 108 | + "#984ea3", |
| 109 | + "#ff7f00", |
| 110 | + "#ffff33", |
| 111 | + "#a65628", |
| 112 | + "#f781bf", |
| 113 | + ] |
| 114 | + plt.ion() |
| 115 | + plt.show(block=False) |
| 116 | + |
| 117 | + def plot_ttc(self): |
| 118 | + self.ttc_ax.cla() |
| 119 | + self.ttc_vel_ax.cla() |
| 120 | + |
| 121 | + n_ttc_data = int(self.ego_ttc_data.layout.dim[1].size) |
| 122 | + ego_ttc_time = self.ego_ttc_data.data[n_ttc_data : 2 * n_ttc_data] |
| 123 | + ego_ttc_dist = self.ego_ttc_data.data[2 * n_ttc_data : 3 * n_ttc_data] |
| 124 | + |
| 125 | + self.ttc_ax.grid() |
| 126 | + self.ttc_ax.set_xlabel("ego time") |
| 127 | + self.ttc_ax.set_ylabel("ego dist") |
| 128 | + time_dist_plot = self.ttc_ax.plot(ego_ttc_time, ego_ttc_dist, label="time-dist", c="orange") |
| 129 | + self.ttc_ax.set_xlim(min(ego_ttc_time) - 2.0, max(ego_ttc_time) + 3.0) |
| 130 | + self.ttc_ax.set_ylim(min(ego_ttc_dist) - 2.0, max(ego_ttc_dist) + 3.0) |
| 131 | + for npc, color in zip(self.npc_vehicles, cycle(self.color_list)): |
| 132 | + t0, t1 = npc.collision_start_time, npc.collision_end_time |
| 133 | + d0, d1 = npc.collision_start_dist, npc.collision_end_dist |
| 134 | + self.ttc_ax.fill( |
| 135 | + [t0, t0, t1, t1, 0, 0], |
| 136 | + [d0, 0, 0, d1, d1, d0], |
| 137 | + c=color, |
| 138 | + alpha=0.2, |
| 139 | + ) |
| 140 | + |
| 141 | + dd = [d1 - d0 for d0, d1 in zip(ego_ttc_dist, ego_ttc_dist[1:])] |
| 142 | + dt = [t1 - t0 for t0, t1 in zip(ego_ttc_time, ego_ttc_time[1:])] |
| 143 | + v = [d / t for d, t in zip(dd, dt)] |
| 144 | + self.ttc_vel_ax.yaxis.set_label_position("right") |
| 145 | + self.ttc_vel_ax.set_ylabel("ego velocity") |
| 146 | + self.ttc_vel_ax.set_ylim(0.0, max(v) + 1.0) |
| 147 | + time_velocity_plot = self.ttc_vel_ax.plot(ego_ttc_time[1:], v, label="time-v", c="red") |
| 148 | + |
| 149 | + lines = time_dist_plot + time_velocity_plot |
| 150 | + labels = [line.get_label() for line in lines] |
| 151 | + self.ttc_ax.legend(lines, labels, loc="upper left") |
| 152 | + |
| 153 | + def plot_world(self): |
| 154 | + detect_range = self.args.range |
| 155 | + self.world_ax.cla() |
| 156 | + n_ttc_data = int(self.ego_ttc_data.layout.dim[1].size) |
| 157 | + ego_path_x = self.ego_ttc_data.data[3 * n_ttc_data : 4 * n_ttc_data] |
| 158 | + ego_path_y = self.ego_ttc_data.data[4 * n_ttc_data : 5 * n_ttc_data] |
| 159 | + self.world_ax.set_aspect("equal") |
| 160 | + self.world_ax.scatter(ego_path_x[0], ego_path_y[0], marker="x", c="red", s=15) |
| 161 | + min_x, max_x = min(ego_path_x), max(ego_path_x) |
| 162 | + min_y, max_y = min(ego_path_y), max(ego_path_y) |
| 163 | + x_dir = 1 if ego_path_x[-1] > ego_path_x[0] else -1 |
| 164 | + y_dir = 1 if ego_path_y[-1] > ego_path_y[0] else -1 |
| 165 | + min_x = min_x - detect_range if x_dir == 1 else min_x - 20 |
| 166 | + max_x = max_x + detect_range if x_dir == -1 else max_x + 20 |
| 167 | + min_y = min_y - detect_range if y_dir == 1 else min_y - 20 |
| 168 | + max_y = max_y + detect_range if y_dir == -1 else max_y + 20 |
| 169 | + self.world_ax.set_xlim(min_x, max_x) |
| 170 | + self.world_ax.set_ylim(min_y, max_y) |
| 171 | + arrows = [ |
| 172 | + (x0, y0, math.atan2(x1 - x0, y1 - y0)) |
| 173 | + for (x0, y0, x1, y1) in zip( |
| 174 | + ego_path_x[0:-1:5], |
| 175 | + ego_path_y[0:-1:5], |
| 176 | + ego_path_x[4:-1:5], |
| 177 | + ego_path_y[4:-1:5], |
| 178 | + ) |
| 179 | + ] |
| 180 | + for x, y, th in arrows: |
| 181 | + self.world_ax.arrow( |
| 182 | + x, |
| 183 | + y, |
| 184 | + math.sin(th) * 0.5, |
| 185 | + math.cos(th) * 0.5, |
| 186 | + head_width=0.1, |
| 187 | + head_length=0.1, |
| 188 | + ) |
| 189 | + |
| 190 | + for npc, color in zip(self.npc_vehicles, cycle(self.color_list)): |
| 191 | + x, y, th, w, h = npc.x, npc.y, npc.th, npc.width, npc.height |
| 192 | + bbox = np.array( |
| 193 | + [ |
| 194 | + [-w / 2, -h / 2], |
| 195 | + [+w / 2, -h / 2], |
| 196 | + [+w / 2, +h / 2], |
| 197 | + [-w / 2, +h / 2], |
| 198 | + [-w / 2, -h / 2], |
| 199 | + ] |
| 200 | + ).transpose() |
| 201 | + Rth = np.array([[math.cos(th), -math.sin(th)], [math.sin(th), math.cos(th)]]) |
| 202 | + bbox_rot = Rth @ bbox |
| 203 | + self.world_ax.fill(bbox_rot[0, :] + x, bbox_rot[1, :] + y, color, alpha=0.5) |
| 204 | + self.world_ax.plot( |
| 205 | + [npc.first_collision_x, npc.last_collision_x], |
| 206 | + [npc.first_collision_y, npc.last_collision_y], |
| 207 | + c=color, |
| 208 | + linewidth=3.0, |
| 209 | + ) |
| 210 | + if npc.dangerous: |
| 211 | + self.world_ax.plot(npc.pred_x, npc.pred_y, c=color, linewidth=1.5) |
| 212 | + else: |
| 213 | + self.world_ax.plot(npc.pred_x, npc.pred_y, c=color, linestyle="--") |
| 214 | + |
| 215 | + self.world_ax.plot(ego_path_x, ego_path_y, c="k", linestyle="--") |
| 216 | + |
| 217 | + def cleanup(self): |
| 218 | + if self.args.save: |
| 219 | + kwargs_write = {"fps": self.args.fps, "quantizer": "nq"} |
| 220 | + imageio.mimsave("./" + self.args.gif + ".gif", self.images, **kwargs_write) |
| 221 | + |
| 222 | + def on_plot_timer(self): |
| 223 | + with self.lock: |
| 224 | + if (not self.ego_ttc_data) or (not self.object_ttc_data): |
| 225 | + return |
| 226 | + |
| 227 | + if not self.last_sub: |
| 228 | + return |
| 229 | + |
| 230 | + now = time.time() |
| 231 | + if (now - self.last_sub) > 1.0: |
| 232 | + print("elapsed more than 1sec from last sub, exit/save fig") |
| 233 | + self.cleanup() |
| 234 | + |
| 235 | + self.plot_ttc() |
| 236 | + self.plot_world() |
| 237 | + self.fig.canvas.flush_events() |
| 238 | + |
| 239 | + if self.args.save: |
| 240 | + image = np.frombuffer(self.fig.canvas.tostring_rgb(), dtype="uint8") |
| 241 | + image = image.reshape(self.fig.canvas.get_width_height()[::-1] + (3,)) |
| 242 | + self.images.append(image) |
| 243 | + |
| 244 | + def on_ego_ttc(self, msg): |
| 245 | + with self.lock: |
| 246 | + if int(msg.data[0]) == self.lane_id: |
| 247 | + self.ego_ttc_data = msg |
| 248 | + self.last_sub = time.time() |
| 249 | + |
| 250 | + def parse_npc_vehicles(self): |
| 251 | + self.npc_vehicles = [] |
| 252 | + n_npc_vehicles = int(self.object_ttc_data.layout.dim[0].size) |
| 253 | + npc_data_size = int(self.object_ttc_data.layout.dim[1].size) |
| 254 | + for i in range(1, n_npc_vehicles): |
| 255 | + data = self.object_ttc_data.data[i * npc_data_size : (i + 1) * npc_data_size] |
| 256 | + self.npc_vehicles.append(NPC(data)) |
| 257 | + |
| 258 | + def on_object_ttc(self, msg): |
| 259 | + with self.lock: |
| 260 | + if int(msg.data[0]) == self.lane_id: |
| 261 | + self.object_ttc_data = msg |
| 262 | + self.parse_npc_vehicles() |
| 263 | + self.last_sub = time.time() |
| 264 | + |
| 265 | + |
| 266 | +if __name__ == "__main__": |
| 267 | + parser = argparse.ArgumentParser() |
| 268 | + parser.add_argument( |
| 269 | + "--lane_id", |
| 270 | + type=int, |
| 271 | + required=True, |
| 272 | + help="lane_id to analyze", |
| 273 | + ) |
| 274 | + parser.add_argument( |
| 275 | + "--range", |
| 276 | + type=float, |
| 277 | + default=60, |
| 278 | + help="detect range for drawing", |
| 279 | + ) |
| 280 | + parser.add_argument("-s", "--save", action="store_true", help="flag to save gif") |
| 281 | + parser.add_argument("--gif", type=str, default="ttc", help="filename of gif file") |
| 282 | + parser.add_argument("--fps", type=float, default=5, help="fps of gif") |
| 283 | + args = parser.parse_args() |
| 284 | + |
| 285 | + rclpy.init() |
| 286 | + visualizer = TTCVisualizer(args) |
| 287 | + rclpy.spin(visualizer) |
0 commit comments