From 14f1386fc0e377d5880de5c3b99ab7cff93af3fa Mon Sep 17 00:00:00 2001 From: Levi Burner Date: Wed, 21 Jun 2023 17:49:18 -0400 Subject: [PATCH] renderer.py improve accuracy and support depth precision and mapping flags By more carefully inverting the operations performed when OpenGL transforms metric distance to window coordinates, more accurate depth values can be estimated from the depth buffer --- python/mujoco/renderer.py | 45 +++++++++++++++++++++++++++++++++------ 1 file changed, 38 insertions(+), 7 deletions(-) diff --git a/python/mujoco/renderer.py b/python/mujoco/renderer.py index e96cd89e..e83f2435 100644 --- a/python/mujoco/renderer.py +++ b/python/mujoco/renderer.py @@ -32,7 +32,9 @@ class Renderer: model: _structs.MjModel, height: int = 240, width: int = 320, - max_geom: int = 10000 + max_geom: int = 10000, + depth_mapping: _enums.mjtDepthMapping = _enums.mjtDepthMapping.mjDB_NEGONETOONE, + depth_precision: _enums.mjtDepthPrecision = _enums.mjtDepthPrecision.mjDB_INT24, ) -> None: """Initializes a new `Renderer`. @@ -43,6 +45,8 @@ class Renderer: max_geom: Optional integer specifying the maximum number of geoms that can be rendered in the same scene. If None this will be chosen automatically based on the estimated maximum number of renderable geoms in the model. + depth_mapping: Type of mapping from znear to zfar to z buffer ndc values + depth_precision: Precision of z buffer Raises: ValueError: If `camera_id` is outside the valid range, or if `width` or `height` exceed the dimensions of MuJoCo's offscreen framebuffer. @@ -70,6 +74,7 @@ the clause: self._width = width self._height = height self._model = model + self._depth_mapping = depth_mapping self._scene = _structs.MjvScene(model=model, maxgeom=max_geom) self._scene_option = _structs.MjvOption() @@ -81,7 +86,7 @@ the clause: self._gl_context = gl_context.GLContext(width, height) # type: ignore self._gl_context.make_current() self._mjr_context = _render.MjrContext( - model, _enums.mjtFontScale.mjFONTSCALE_150.value + model, _enums.mjtFontScale.mjFONTSCALE_150.value, depth_mapping, depth_precision ) _render.mjr_setBuffer( _enums.mjtFramebuffer.mjFB_OFFSCREEN.value, self._mjr_context @@ -138,7 +143,8 @@ the clause: """ original_flags = self._scene.flags.copy() - if self._segmentation_rendering: + # Using segmented rendering for depth makes the calculated depth more accurate at far distances + if self._depth_rendering or self._segmentation_rendering: self._scene.flags[_enums.mjtRndFlag.mjRND_SEGMENT] = True self._scene.flags[_enums.mjtRndFlag.mjRND_IDCOLOR] = True @@ -173,11 +179,36 @@ the clause: near = self._model.vis.map.znear * extent far = self._model.vis.map.zfar * extent - # Convert from [0 1] to depth in units of length, see links below: - # http://stackoverflow.com/a/6657284/1461210 - # https://www.khronos.org/opengl/wiki/Depth_Buffer_Precision - out = near / (1 - out * (1 - near / far)) + # Calculate OpenGL perspective matrix values in float32 precision + # so they are close to what glFrustum returns + # https://registry.khronos.org/OpenGL-Refpages/gl2.1/xhtml/glFrustum.xml + zfar = np.float32(far) + znear = np.float32(near) + C = -(zfar + znear)/(zfar - znear) + D = -(np.float32(2)*zfar*znear)/(zfar - znear) + # In reverse Z mode the perspective matrix is transformed by the following + if self._depth_mapping == _enums.mjtDepthMapping.mjDB_ONETOZERO: + C = np.float32(-0.5)*C - np.float32(0.5) + D = np.float32(-0.5)*D + + # We need 64 bits to convert Z from ndc to metric depth without noticeable losses in precision + out_64 = out.astype(np.float64) + + # Convert depth from window coordinates to normalized device coordinates + # In reversed Z mode the mapping is identity + # https://registry.khronos.org/OpenGL-Refpages/gl4/html/glDepthRange.xhtml + if self._depth_mapping == _enums.mjtDepthMapping.mjDB_NEGONETOONE: + out_64 = 2.0*out_64 - 1.0 + + # Undo OpenGL projection + out_64 = D / (out_64 + C) + + # Cast result back to float32 for backwards compatibility + out[:] = out_64.astype(np.float32) + + # Reset scene flags. + np.copyto(self._scene.flags, original_flags) elif self._segmentation_rendering: _render.mjr_readPixels(out, None, self._rect, self._mjr_context)