1use crate::wrappers::mj_visualization::MjvScene;
3use crate::wrappers::mj_rendering::MjrContext;
4use crate::builder_setters;
5use crate::prelude::*;
6
7use glfw_mjrc_fork as glfw;
8use glfw_mjrc_fork::{Context, InitError, PWindow, WindowHint};
9
10use bitflags::bitflags;
11use png::Encoder;
12
13use std::io::{self, BufWriter, ErrorKind, Write};
14use std::fmt::Display;
15use std::error::Error;
16use std::path::Path;
17use std::fs::File;
18
19const RGB_NOT_FOUND_ERR_STR: &str = "RGB rendering is not enabled (renderer.with_rgb_rendering(true))";
20const DEPTH_NOT_FOUND_ERR_STR: &str = "depth rendering is not enabled (renderer.with_depth_rendering(true))";
21const INVALID_INPUT_SIZE: &str = "the input width and height don't match the renderer's configuration";
22const EXTRA_INTERNAL_VISUAL_GEOMS: usize = 100;
23
24
25#[derive(Debug)]
27pub struct MjRendererBuilder {
28 width: u32,
29 height: u32,
30 num_visual_internal_geom: u32,
31 num_visual_user_geom: u32,
32 rgb: bool,
33 depth: bool,
34 font_scale: MjtFontScale,
35 camera: MjvCamera,
36 opts: MjvOption,
37}
38
39impl MjRendererBuilder {
40 pub fn new() -> Self {
48 Self {
49 width: 0, height: 0,
50 num_visual_internal_geom: EXTRA_INTERNAL_VISUAL_GEOMS as u32, num_visual_user_geom: 0,
51 rgb: true, depth: false, font_scale: MjtFontScale::mjFONTSCALE_100,
52 camera: MjvCamera::default(), opts: MjvOption::default()
53 }
54 }
55
56 builder_setters! {
57 width: u32; "
58image width.
59
60<div class=\"warning\">
61
62The width must be less or equal to the offscreen buffer width,
63which can be configured at the top of the model's XML like so:
64
65```xml
66<visual>
67 <global offwidth=\"1920\" .../>
68</visual>
69```
70
71</div>";
72
73 height: u32; "\
74image height.
75
76<div class=\"warning\">
77
78The height must be less or equal to the offscreen buffer height,
79which can be configured at the top of the model's XML like so:
80
81```xml
82<visual>
83 <global offheight=\"1080\" .../>
84</visual>
85```
86
87</div>";
88
89 num_visual_internal_geom: u32; "\
90 maximum number of additional visual-only internal geoms to allocate for.
91 Note that the total number of geoms in the internal scene will be
92 `num_visual_internal_geom` + `num_visual_user_geom`.";
93
94 num_visual_user_geom: u32; "maximum number of additional visual-only user geoms (drawn by the user).";
95 rgb: bool; "RGB rendering enabled (true) or disabled (false).";
96 depth: bool; "depth rendering enabled (true) or disabled (false).";
97 font_scale: MjtFontScale; "font scale of drawn text (with [MjrContext]).";
98 camera: MjvCamera; "camera used for drawing.";
99 opts: MjvOption; "visualization options.";
100
101 }
102
103 pub fn build(self, model: &MjModel) -> Result<MjRenderer<'_>, RendererError> {
105 let mut height = self.height;
107 let mut width = self.width;
108 if width == 0 && height == 0 {
109 let global = &model.vis().global;
110 height = global.offheight as u32;
111 width = global.offwidth as u32;
112 }
113
114 let mut glfw = glfw::init_no_callbacks()
115 .map_err(|err| RendererError::GlfwInitError(err))?;
116
117 glfw.window_hint(WindowHint::Visible(false));
119 let (mut window, _) = match glfw.create_window(
120 width, height,
121 "", glfw::WindowMode::Windowed
122 ) {
123 Some(x) => Ok(x),
124 None => Err(RendererError::WindowCreationError)
125 }?;
126
127 window.make_current();
128 glfw.set_swap_interval(glfw::SwapInterval::None);
129
130 let mut context = MjrContext::new(model);
132 context.offscreen();
133
134 let scene = MjvScene::new(
136 model,
137 model.ffi().ngeom as usize + self.num_visual_internal_geom as usize
138 + self.num_visual_user_geom as usize
139 );
140
141 let user_scene = MjvScene::new(
142 model,
143 self.num_visual_user_geom as usize
144 );
145
146 let renderer = MjRenderer {
148 scene, user_scene, context, window, model, camera: self.camera, option: self.opts,
149 flags: RendererFlags::empty(), rgb: None, depth: None,
150 width: width as usize, height: height as usize
151 } .with_rgb_rendering(self.rgb)
153 .with_depth_rendering(self.depth);
154
155 Ok(renderer)
156 }
157}
158
159
160impl Default for MjRendererBuilder {
161 fn default() -> Self {
162 Self::new()
163 }
164}
165
166pub struct MjRenderer<'m> {
169 scene: MjvScene<'m>,
170 user_scene: MjvScene<'m>,
171 context: MjrContext,
172 model: &'m MjModel,
173
174 window: PWindow,
176
177 camera: MjvCamera,
179 option: MjvOption,
180 flags: RendererFlags,
181
182 rgb: Option<Box<[u8]>>,
186 depth: Option<Box<[f32]>>,
187
188 width: usize,
189 height: usize,
190}
191
192impl<'m> MjRenderer<'m> {
193 pub fn new(model: &'m MjModel, width: usize, height: usize, max_geom: usize) -> Result<Self, RendererError> {
221 let mut glfw = glfw::init_no_callbacks()
222 .map_err(|err| RendererError::GlfwInitError(err))?;
223
224 glfw.window_hint(WindowHint::Visible(false));
226 let (mut _window, _) = match glfw.create_window(width as u32, height as u32, "", glfw::WindowMode::Windowed) {
227 Some(x) => Ok(x),
228 None => Err(RendererError::WindowCreationError)
229 }?;
230
231 _window.make_current();
232 glfw.set_swap_interval(glfw::SwapInterval::None);
233
234 let mut context = MjrContext::new(model);
236 context.offscreen();
237
238 let scene = MjvScene::new(model, model.ffi().ngeom as usize + max_geom + EXTRA_INTERNAL_VISUAL_GEOMS);
240 let user_scene = MjvScene::new(model, max_geom);
241
242 let camera = MjvCamera::new_free(model);
243 let option = MjvOption::default();
244
245 let mut s = Self {
246 scene, user_scene, context, window: _window, model, camera, option,
247 flags: RendererFlags::empty(), rgb: None, depth: None,
248 width, height
249 };
250
251 s = s.with_rgb_rendering(true);
252 Ok(s)
253 }
254
255 pub fn builder() -> MjRendererBuilder {
257 MjRendererBuilder::new()
258 }
259
260 pub fn scene(&self) -> &MjvScene<'m>{
262 &self.scene
263 }
264
265 pub fn user_scene(&self) -> &MjvScene<'m>{
267 &self.user_scene
268 }
269
270 pub fn user_scene_mut(&mut self) -> &mut MjvScene<'m>{
272 &mut self.user_scene
273 }
274
275 pub fn opts(&self) -> &MjvOption {
277 &self.option
278 }
279
280 pub fn opts_mut(&mut self) -> &mut MjvOption {
282 &mut self.option
283 }
284
285 pub fn camera(&self) -> &MjvCamera {
287 &self.camera
288 }
289
290 pub fn camera_mut(&mut self) -> &mut MjvCamera {
292 &mut self.camera
293 }
294
295 pub fn rgb_enabled(&self) -> bool {
297 self.flags.contains(RendererFlags::RENDER_RGB)
298 }
299
300 pub fn depth_enabled(&self) -> bool {
302 self.flags.contains(RendererFlags::RENDER_DEPTH)
303 }
304
305 pub fn set_font_scale(&mut self, font_scale: MjtFontScale) {
307 self.context.change_font(font_scale);
308 }
309
310 pub fn set_opts(&mut self, options: MjvOption) {
312 self.option = options;
313 }
314
315 pub fn set_camera(&mut self, camera: MjvCamera) {
317 self.camera = camera;
318 }
319
320 pub fn set_rgb_rendering(&mut self, enable: bool) {
322 self.flags.set(RendererFlags::RENDER_RGB, enable);
323 self.rgb = if enable { Some(vec![0; 3 * self.width * self.height].into_boxed_slice()) } else { None } ;
324 }
325
326 pub fn set_depth_rendering(&mut self, enable: bool) {
328 self.flags.set(RendererFlags::RENDER_DEPTH, enable);
329 self.depth = if enable { Some(vec![0.0; self.width * self.height].into_boxed_slice()) } else { None } ;
330 }
331
332 pub fn with_font_scale(mut self, font_scale: MjtFontScale) -> Self {
334 self.set_font_scale(font_scale);
335 self
336 }
337
338 pub fn with_opts(mut self, options: MjvOption) -> Self {
340 self.set_opts(options);
341 self
342 }
343
344 pub fn with_camera(mut self, camera: MjvCamera) -> Self {
346 self.set_camera(camera);
347 self
348 }
349
350 pub fn with_rgb_rendering(mut self, enable: bool) -> Self {
352 self.set_rgb_rendering(enable);
353 self
354 }
355
356 pub fn with_depth_rendering(mut self, enable: bool) -> Self {
358 self.set_depth_rendering(enable);
359 self
360 }
361
362 pub fn sync(&mut self, data: &mut MjData) {
364 let model_data_ptr = unsafe { data.model().__raw() };
365 let bound_model_ptr = unsafe { self.model.__raw() };
366 assert_eq!(model_data_ptr, bound_model_ptr, "'data' must be created from the same model as the renderer.");
367
368 self.scene.update(data, &self.option, &MjvPerturb::default(), &mut self.camera);
369
370 sync_geoms(&self.user_scene, &mut self.scene)
372 .expect("could not sync the user scene with the internal scene; this is a bug, please report it.");
373
374 self.render();
375 }
376
377 pub fn rgb_flat(&self) -> Option<&[u8]> {
379 self.rgb.as_deref()
380 }
381
382 pub fn rgb<const WIDTH: usize, const HEIGHT: usize>(&self) -> io::Result<&[[[u8; 3]; WIDTH]; HEIGHT]> {
385 if let Some(flat) = self.rgb_flat() {
386 if flat.len() == WIDTH * HEIGHT * 3 {
387 let p_shaped = flat.as_ptr() as *const [[[u8; 3]; WIDTH]; HEIGHT];
388
389 Ok(unsafe { p_shaped.as_ref().unwrap() })
393 }
394 else {
395 Err(io::Error::new(io::ErrorKind::InvalidInput, INVALID_INPUT_SIZE))
396 }
397 }
398 else {
399 Err(io::Error::new(io::ErrorKind::NotFound, RGB_NOT_FOUND_ERR_STR))
400 }
401 }
402
403 pub fn depth_flat(&self) -> Option<&[f32]> {
405 self.depth.as_deref()
406 }
407
408 pub fn depth<const WIDTH: usize, const HEIGHT: usize>(&self) -> io::Result<&[[f32; WIDTH]; HEIGHT]> {
411 if let Some(flat) = self.depth_flat() {
412 if flat.len() == WIDTH * HEIGHT {
413 let p_shaped = flat.as_ptr() as *const [[f32; WIDTH]; HEIGHT];
414
415 Ok(unsafe { p_shaped.as_ref().unwrap() })
419 }
420 else {
421 Err(io::Error::new(io::ErrorKind::InvalidInput, INVALID_INPUT_SIZE))
422 }
423 }
424 else {
425 Err(io::Error::new(io::ErrorKind::NotFound, DEPTH_NOT_FOUND_ERR_STR))
426 }
427 }
428
429 pub fn save_rgb<T: AsRef<Path>>(&self, path: T) -> io::Result<()> {
434 if let Some(rgb) = &self.rgb {
435 let file = File::create(path.as_ref())?;
436 let w = BufWriter::new(file);
437
438 let mut encoder = Encoder::new(w, self.width as u32, self.height as u32);
439 encoder.set_color(png::ColorType::Rgb);
440 encoder.set_depth(png::BitDepth::Eight);
441 encoder.set_compression(png::Compression::NoCompression);
442
443 let mut writer = encoder.write_header()?;
444 writer.write_image_data(rgb)?;
445 Ok(())
446 }
447 else {
448 Err(io::Error::new(ErrorKind::NotFound, RGB_NOT_FOUND_ERR_STR))
449 }
450 }
451
452 pub fn save_depth<T: AsRef<Path>>(&self, path: T, normalize: bool) -> io::Result<(f32, f32)> {
464 if let Some(depth) = &self.depth {
465 let file = File::create(path.as_ref())?;
466 let w = BufWriter::new(file);
467
468 let mut encoder = Encoder::new(w, self.width as u32, self.height as u32);
469 encoder.set_color(png::ColorType::Grayscale);
470 encoder.set_depth(png::BitDepth::Sixteen);
471 encoder.set_compression(png::Compression::NoCompression);
472
473 let (norm, min, max) =
474 if normalize {
475 let max = depth.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
476 let min = depth.iter().cloned().fold(f32::INFINITY, f32::min);
477 (depth.iter().flat_map(|&x| (((x - min) / (max - min) * 65535.0).min(65535.0) as u16).to_be_bytes()).collect::<Box<_>>(), min, max)
478 }
479 else {
480 (depth.iter().flat_map(|&x| ((x * 65535.0).min(65535.0) as u16).to_be_bytes()).collect::<Box<_>>(), 0.0, 1.0)
481 };
482
483 let mut writer = encoder.write_header()?;
484 writer.write_image_data(&norm)?;
485 Ok((min, max))
486 }
487 else {
488 Err(io::Error::new(ErrorKind::NotFound, DEPTH_NOT_FOUND_ERR_STR))
489 }
490 }
491
492 pub fn save_depth_raw<T: AsRef<Path>>(&self, path: T) -> io::Result<()> {
499 if let Some(depth) = &self.depth {
500 let file = File::create(path.as_ref())?;
501 let mut writer = BufWriter::new(file);
502
503 let p = unsafe { std::slice::from_raw_parts(
505 depth.as_ptr() as *const u8,
506 std::mem::size_of::<f32>() * depth.len()
507 ) };
508
509 writer.write_all(p)?;
510 Ok(())
511 }
512 else {
513 Err(io::Error::new(ErrorKind::NotFound, DEPTH_NOT_FOUND_ERR_STR))
514 }
515 }
516
517 fn render(&mut self) {
520 self.window.make_current();
521 let vp = MjrRectangle::new(0, 0, self.width as i32, self.height as i32);
522 self.scene.render(&vp, &self.context);
523
524 let flat_rgb = self.rgb.as_deref_mut();
526 let flat_depth = self.depth.as_deref_mut();
527
528 self.context.read_pixels(
530 flat_rgb,
531 flat_depth,
532 &vp
533 );
534
535 if let Some(depth) = self.depth.as_deref_mut() {
537 let map = &self.model.vis().map;
538 let near = map.znear;
539 let far = map.zfar;
540 for value in depth {
541 let z_ndc = 2.0 * *value - 1.0;
542 *value = 2.0 * near * far / (far + near - z_ndc * (far - near));
543 }
544 }
545 }
546}
547
548
549#[derive(Debug)]
550pub enum RendererError {
551 GlfwInitError(InitError),
552 WindowCreationError,
553}
554
555impl Display for RendererError {
556 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
557 match self {
558 Self::GlfwInitError(e) => write!(f, "glfw failed to initialize: {}", e),
559 Self::WindowCreationError => write!(f, "failed to create window"),
560 }
561 }
562}
563
564impl Error for RendererError {
565 fn source(&self) -> Option<&(dyn Error + 'static)> {
566 match self {
567 Self::GlfwInitError(e) => Some(e),
568 _ => None
569 }
570 }
571}
572
573bitflags! {
574 struct RendererFlags: u8 {
576 const RENDER_RGB = 1 << 0;
577 const RENDER_DEPTH = 1 << 1;
578 }
579}
580
581
582
583