1use gnuplot::{AutoOption, AxesCommon, Caption, Color, Figure, LineWidth, PointSize, PointSymbol};
6use rust_robotics_core::{Obstacles, Path2D, Point2D, Pose2D};
7
8pub mod colors {
10 pub const BLACK: &str = "#000000";
11 pub const RED: &str = "#FF0000";
12 pub const GREEN: &str = "#00FF00";
13 pub const BLUE: &str = "#0000FF";
14 pub const YELLOW: &str = "#FFFF00";
15 pub const CYAN: &str = "#00FFFF";
16 pub const MAGENTA: &str = "#FF00FF";
17 pub const ORANGE: &str = "#FFA500";
18 pub const PURPLE: &str = "#800080";
19 pub const GRAY: &str = "#808080";
20
21 pub const OBSTACLE: &str = BLACK;
23 pub const START: &str = GREEN;
24 pub const GOAL: &str = BLUE;
25 pub const PATH: &str = RED;
26 pub const ROBOT: &str = CYAN;
27 pub const ESTIMATED: &str = "#35C788";
28 pub const GROUND_TRUTH: &str = BLUE;
29 pub const MEASUREMENT: &str = "#DD3355";
30 pub const DEAD_RECKONING: &str = YELLOW;
31}
32
33#[derive(Debug, Clone)]
35pub struct PathStyle {
36 pub color: String,
37 pub line_width: f64,
38 pub caption: String,
39}
40
41impl PathStyle {
42 pub fn new(color: &str, caption: &str) -> Self {
43 Self {
44 color: color.to_string(),
45 line_width: 2.0,
46 caption: caption.to_string(),
47 }
48 }
49
50 pub fn with_line_width(mut self, width: f64) -> Self {
51 self.line_width = width;
52 self
53 }
54}
55
56impl Default for PathStyle {
57 fn default() -> Self {
58 Self {
59 color: colors::PATH.to_string(),
60 line_width: 2.0,
61 caption: "Path".to_string(),
62 }
63 }
64}
65
66#[derive(Debug, Clone)]
68pub struct PointStyle {
69 pub color: String,
70 pub size: f64,
71 pub symbol: char,
72 pub caption: String,
73}
74
75impl PointStyle {
76 pub fn new(color: &str, caption: &str) -> Self {
77 Self {
78 color: color.to_string(),
79 size: 1.0,
80 symbol: 'O',
81 caption: caption.to_string(),
82 }
83 }
84
85 pub fn with_size(mut self, size: f64) -> Self {
86 self.size = size;
87 self
88 }
89
90 pub fn with_symbol(mut self, symbol: char) -> Self {
91 self.symbol = symbol;
92 self
93 }
94}
95
96enum PlotCommand {
98 Path {
99 x: Vec<f64>,
100 y: Vec<f64>,
101 style: PathStyle,
102 },
103 Points {
104 x: Vec<f64>,
105 y: Vec<f64>,
106 style: PointStyle,
107 },
108}
109
110pub struct Visualizer {
112 figure: Figure,
113 title: String,
114 x_label: String,
115 y_label: String,
116 x_range: Option<(f64, f64)>,
117 y_range: Option<(f64, f64)>,
118 aspect_ratio: Option<f64>,
119 commands: Vec<PlotCommand>,
120}
121
122impl Visualizer {
123 pub fn new() -> Self {
125 Self {
126 figure: Figure::new(),
127 title: String::new(),
128 x_label: "X [m]".to_string(),
129 y_label: "Y [m]".to_string(),
130 x_range: None,
131 y_range: None,
132 aspect_ratio: Some(1.0),
133 commands: Vec::new(),
134 }
135 }
136
137 pub fn set_title(&mut self, title: &str) -> &mut Self {
139 self.title = title.to_string();
140 self
141 }
142
143 pub fn set_x_label(&mut self, label: &str) -> &mut Self {
145 self.x_label = label.to_string();
146 self
147 }
148
149 pub fn set_y_label(&mut self, label: &str) -> &mut Self {
151 self.y_label = label.to_string();
152 self
153 }
154
155 pub fn set_x_range(&mut self, min: f64, max: f64) -> &mut Self {
157 self.x_range = Some((min, max));
158 self
159 }
160
161 pub fn set_y_range(&mut self, min: f64, max: f64) -> &mut Self {
163 self.y_range = Some((min, max));
164 self
165 }
166
167 pub fn set_aspect_ratio(&mut self, ratio: Option<f64>) -> &mut Self {
169 self.aspect_ratio = ratio;
170 self
171 }
172
173 pub fn figure_mut(&mut self) -> &mut Figure {
175 &mut self.figure
176 }
177
178 pub fn plot_path(&mut self, path: &Path2D, style: &PathStyle) -> &mut Self {
180 let x: Vec<f64> = path.points.iter().map(|p| p.x).collect();
181 let y: Vec<f64> = path.points.iter().map(|p| p.y).collect();
182 self.commands.push(PlotCommand::Path {
183 x,
184 y,
185 style: style.clone(),
186 });
187 self
188 }
189
190 pub fn plot_path_xy(&mut self, x: &[f64], y: &[f64], style: &PathStyle) -> &mut Self {
192 self.commands.push(PlotCommand::Path {
193 x: x.to_vec(),
194 y: y.to_vec(),
195 style: style.clone(),
196 });
197 self
198 }
199
200 pub fn plot_obstacles(&mut self, obstacles: &Obstacles) -> &mut Self {
202 let x: Vec<f64> = obstacles.points.iter().map(|p| p.x).collect();
203 let y: Vec<f64> = obstacles.points.iter().map(|p| p.y).collect();
204 self.commands.push(PlotCommand::Points {
205 x,
206 y,
207 style: PointStyle {
208 color: colors::OBSTACLE.to_string(),
209 size: 0.5,
210 symbol: 'S',
211 caption: "Obstacles".to_string(),
212 },
213 });
214 self
215 }
216
217 pub fn plot_obstacles_xy(&mut self, ox: &[f64], oy: &[f64]) -> &mut Self {
219 self.commands.push(PlotCommand::Points {
220 x: ox.to_vec(),
221 y: oy.to_vec(),
222 style: PointStyle {
223 color: colors::OBSTACLE.to_string(),
224 size: 0.5,
225 symbol: 'S',
226 caption: "Obstacles".to_string(),
227 },
228 });
229 self
230 }
231
232 pub fn plot_point(&mut self, point: Point2D, style: &PointStyle) -> &mut Self {
234 self.commands.push(PlotCommand::Points {
235 x: vec![point.x],
236 y: vec![point.y],
237 style: style.clone(),
238 });
239 self
240 }
241
242 pub fn plot_points(&mut self, points: &[Point2D], style: &PointStyle) -> &mut Self {
244 let x: Vec<f64> = points.iter().map(|p| p.x).collect();
245 let y: Vec<f64> = points.iter().map(|p| p.y).collect();
246 self.commands.push(PlotCommand::Points {
247 x,
248 y,
249 style: style.clone(),
250 });
251 self
252 }
253
254 pub fn plot_points_xy(&mut self, x: &[f64], y: &[f64], style: &PointStyle) -> &mut Self {
256 self.commands.push(PlotCommand::Points {
257 x: x.to_vec(),
258 y: y.to_vec(),
259 style: style.clone(),
260 });
261 self
262 }
263
264 pub fn plot_robot(&mut self, pose: &Pose2D, size: f64) -> &mut Self {
266 self.commands.push(PlotCommand::Points {
268 x: vec![pose.x],
269 y: vec![pose.y],
270 style: PointStyle {
271 color: colors::ROBOT.to_string(),
272 size,
273 symbol: 'O',
274 caption: "Robot".to_string(),
275 },
276 });
277
278 let arrow_len = size * 0.5;
280 let end_x = pose.x + arrow_len * pose.yaw.cos();
281 let end_y = pose.y + arrow_len * pose.yaw.sin();
282 self.commands.push(PlotCommand::Path {
283 x: vec![pose.x, end_x],
284 y: vec![pose.y, end_y],
285 style: PathStyle {
286 color: colors::ROBOT.to_string(),
287 line_width: 2.0,
288 caption: String::new(),
289 },
290 });
291 self
292 }
293
294 pub fn plot_start(&mut self, point: Point2D) -> &mut Self {
296 self.plot_point(
297 point,
298 &PointStyle::new(colors::START, "Start").with_size(1.5),
299 )
300 }
301
302 pub fn plot_goal(&mut self, point: Point2D) -> &mut Self {
304 self.plot_point(point, &PointStyle::new(colors::GOAL, "Goal").with_size(1.5))
305 }
306
307 pub fn show(&mut self) -> Result<(), String> {
309 self.apply_settings();
310 self.figure.show().map_err(|e| e.to_string()).map(|_| ())
311 }
312
313 pub fn save_png(&mut self, path: &str, width: u32, height: u32) -> Result<(), String> {
315 self.apply_settings();
316 self.figure
317 .save_to_png(path, width, height)
318 .map_err(|e| e.to_string())
319 }
320
321 pub fn save_svg(&mut self, path: &str, width: u32, height: u32) -> Result<(), String> {
323 self.apply_settings();
324 self.figure
325 .save_to_svg(path, width, height)
326 .map_err(|e| e.to_string())
327 }
328
329 fn apply_settings(&mut self) {
330 let axes = self.figure.axes2d();
331
332 for cmd in &self.commands {
334 match cmd {
335 PlotCommand::Path { x, y, style } => {
336 axes.lines(
337 x,
338 y,
339 &[
340 Caption(&style.caption),
341 Color(&style.color),
342 LineWidth(style.line_width),
343 ],
344 );
345 }
346 PlotCommand::Points { x, y, style } => {
347 axes.points(
348 x,
349 y,
350 &[
351 Caption(&style.caption),
352 Color(&style.color),
353 PointSymbol(style.symbol),
354 PointSize(style.size),
355 ],
356 );
357 }
358 }
359 }
360
361 if !self.title.is_empty() {
362 axes.set_title(&self.title, &[]);
363 }
364 axes.set_x_label(&self.x_label, &[]);
365 axes.set_y_label(&self.y_label, &[]);
366
367 if let Some((min, max)) = self.x_range {
368 axes.set_x_range(AutoOption::Fix(min), AutoOption::Fix(max));
369 }
370 if let Some((min, max)) = self.y_range {
371 axes.set_y_range(AutoOption::Fix(min), AutoOption::Fix(max));
372 }
373 if let Some(ratio) = self.aspect_ratio {
374 axes.set_aspect_ratio(AutoOption::Fix(ratio));
375 }
376 }
377}
378
379impl Default for Visualizer {
380 fn default() -> Self {
381 Self::new()
382 }
383}
384
385pub fn quick_plot_path(
387 path: &Path2D,
388 obstacles: Option<&Obstacles>,
389 start: Option<Point2D>,
390 goal: Option<Point2D>,
391 title: &str,
392) -> Visualizer {
393 let mut vis = Visualizer::new();
394 vis.set_title(title);
395
396 if let Some(obs) = obstacles {
397 vis.plot_obstacles(obs);
398 }
399 if let Some(s) = start {
400 vis.plot_start(s);
401 }
402 if let Some(g) = goal {
403 vis.plot_goal(g);
404 }
405 vis.plot_path(path, &PathStyle::default());
406
407 vis
408}
409
410#[cfg(test)]
411mod tests {
412 use super::*;
413
414 #[test]
415 fn test_visualizer_creation() {
416 let vis = Visualizer::new();
417 assert!(vis.aspect_ratio.is_some());
418 }
419
420 #[test]
421 fn test_path_style() {
422 let style = PathStyle::new(colors::RED, "Test Path").with_line_width(3.0);
423 assert_eq!(style.line_width, 3.0);
424 assert_eq!(style.color, colors::RED);
425 }
426}