warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/render/render_opengl.py

@@ -13,12 +13,14 @@ from .utils import tab10_color_map
 
 from collections import defaultdict
 from typing import List, Tuple, Union, Optional
+from enum import Enum
 
 import numpy as np
 import ctypes
 
 Mat44 = Union[List[float], List[List[float]], np.ndarray]
 
+
 wp.set_module_options({"enable_backward": False})
 
 shape_vertex_shader = """
@@ -175,12 +177,13 @@ in vec2 TexCoord;
 
 uniform vec3 color1;
 uniform vec3 color2;
+uniform float farPlane;
 
 uniform vec3 sunDirection;
 
 void main()
 {
-    float y = tanh(FragPos.y*0.01)*0.5+0.5;
+    float y = tanh(FragPos.y/farPlane*10.0)*0.5+0.5;
     float height = sqrt(1.0-y);
 
     float s = pow(0.5, 1.0 / 10.0);
@@ -222,6 +225,42 @@ void main() {
 }
 """
 
+frame_depth_fragment_shader = """
+#version 330 core
+in vec2 TexCoord;
+
+out vec4 FragColor;
+
+uniform sampler2D textureSampler;
+
+vec3 bourkeColorMap(float v) {
+    vec3 c = vec3(1.0, 1.0, 1.0);
+
+    v = clamp(v, 0.0, 1.0); // Ensures v is between 0 and 1
+
+    if (v < 0.25) {
+        c.r = 0.0;
+        c.g = 4.0 * v;
+    } else if (v < 0.5) {
+        c.r = 0.0;
+        c.b = 1.0 + 4.0 * (0.25 - v);
+    } else if (v < 0.75) {
+        c.r = 4.0 * (v - 0.5);
+        c.b = 0.0;
+    } else {
+        c.g = 1.0 + 4.0 * (0.75 - v);
+        c.b = 0.0;
+    }
+
+    return c;
+}
+
+void main() {
+    float depth = texture(textureSampler, TexCoord).r;
+    FragColor = vec4(bourkeColorMap(sqrt(1.0 - depth)), 1.0);
+}
+"""
+
 
 @wp.kernel
 def update_vbo_transforms(
@@ -411,7 +450,7 @@ def assemble_gfx_vertices(
 
 
 @wp.kernel
-def copy_frame(
+def copy_rgb_frame(
     input_img: wp.array(dtype=wp.uint8),
     width: int,
     height: int,
@@ -432,7 +471,26 @@ def copy_frame(
 
 
 @wp.kernel
-def copy_frame_tiles(
+def copy_depth_frame(
+    input_img: wp.array(dtype=wp.float32),
+    width: int,
+    height: int,
+    near: float,
+    far: float,
+    # outputs
+    output_img: wp.array(dtype=wp.float32, ndim=3),
+):
+    w, v = wp.tid()
+    pixel = v * width + w
+    # flip vertically (OpenGL coordinates start at bottom)
+    v = height - v - 1
+    d = 2.0 * input_img[pixel] - 1.0
+    d = 2.0 * near * far / ((far - near) * d - near - far)
+    output_img[v, w, 0] = -d
+
+
+@wp.kernel
+def copy_rgb_frame_tiles(
     input_img: wp.array(dtype=wp.uint8),
     positions: wp.array(dtype=int, ndim=2),
     screen_width: int,
@@ -463,7 +521,34 @@ def copy_frame_tiles(
 
 
 @wp.kernel
-def copy_frame_tile(
+def copy_depth_frame_tiles(
+    input_img: wp.array(dtype=wp.float32),
+    positions: wp.array(dtype=int, ndim=2),
+    screen_width: int,
+    screen_height: int,
+    tile_height: int,
+    near: float,
+    far: float,
+    # outputs
+    output_img: wp.array(dtype=wp.float32, ndim=4),
+):
+    tile, x, y = wp.tid()
+    p = positions[tile]
+    qx = x + p[0]
+    qy = y + p[1]
+    pixel = qy * screen_width + qx
+    # flip vertically (OpenGL coordinates start at bottom)
+    y = tile_height - y - 1
+    if qx >= screen_width or qy >= screen_height:
+        output_img[tile, y, x, 0] = far
+        return  # prevent out-of-bounds access
+    d = 2.0 * input_img[pixel] - 1.0
+    d = 2.0 * near * far / ((far - near) * d - near - far)
+    output_img[tile, y, x, 0] = -d
+
+
+@wp.kernel
+def copy_rgb_frame_tile(
     input_img: wp.array(dtype=wp.uint8),
     offset_x: int,
     offset_y: int,
@@ -765,27 +850,35 @@ class OpenGLRenderer:
         title="Warp sim",
         scaling=1.0,
         fps=60,
-        up_axis="y",
+        up_axis="Y",
         screen_width=1024,
         screen_height=768,
-        near_plane=0.01,
-        far_plane=1000.0,
+        near_plane=1.0,
+        far_plane=100.0,
         camera_fov=45.0,
+        camera_pos=(0.0, 2.0, 10.0),
+        camera_front=(0.0, 0.0, -1.0),
+        camera_up=(0.0, 1.0, 0.0),
         background_color=(0.53, 0.8, 0.92),
         draw_grid=True,
         draw_sky=True,
         draw_axis=True,
         show_info=True,
+        render_wireframe=False,
+        render_depth=False,
         axis_scale=1.0,
-        vsync=True,
+        vsync=False,
         headless=False,
+        enable_backface_culling=True,
+        enable_mouse_interaction=True,
+        enable_keyboard_interaction=True,
     ):
         try:
             import pyglet
-        
+
             # disable error checking for performance
-            pyglet.options['debug_gl'] = False
-            
+            pyglet.options["debug_gl"] = False
+
             from pyglet import gl
             from pyglet.math import Vec3 as PyVec3
             from pyglet.graphics.shader import Shader, ShaderProgram
@@ -801,6 +894,9 @@ class OpenGLRenderer:
         self.draw_sky = draw_sky
         self.draw_axis = draw_axis
         self.show_info = show_info
+        self.render_wireframe = render_wireframe
+        self.render_depth = render_depth
+        self.enable_backface_culling = enable_backface_culling
 
         self._device = wp.get_cuda_device()
         self._title = title
@@ -815,19 +911,26 @@ class OpenGLRenderer:
 
         self.screen_width, self.screen_height = self.window.get_framebuffer_size()
 
-        self._camera_pos = PyVec3(0.0, 2.0, 10.0)
-        self._camera_front = PyVec3(0.0, 0.0, -1.0)
-        self._camera_up = PyVec3(0.0, 1.0, 0.0)
+        self.enable_mouse_interaction = enable_mouse_interaction
+        self.enable_keyboard_interaction = enable_keyboard_interaction
+
+        self._camera_pos = PyVec3(*camera_pos)
+        self._camera_front = PyVec3(*camera_front)
+        self._camera_up = PyVec3(*camera_up)
         self._camera_speed = 0.04
         if isinstance(up_axis, int):
             self._camera_axis = up_axis
         else:
-            self._camera_axis = "xyz".index(up_axis.lower())
+            self._camera_axis = "XYZ".index(up_axis.upper())
         self._yaw, self._pitch = -90.0, 0.0
         self._last_x, self._last_y = self.screen_width // 2, self.screen_height // 2
         self._first_mouse = True
         self._left_mouse_pressed = False
         self._keys_pressed = defaultdict(bool)
+        self._key_callbacks = []
+
+        self.render_2d_callbacks = []
+        self.render_3d_callbacks = []
 
         self.update_view_matrix()
         self.update_projection_matrix()
@@ -850,6 +953,7 @@ class OpenGLRenderer:
         self._shape_gl_buffers = {}
         self._shape_instances = defaultdict(list)
         self._instances = {}
+        self._instance_custom_ids = {}
         self._instance_shape = {}
         self._instance_gl_buffers = {}
         self._instance_transform_gl_buffer = None
@@ -858,6 +962,8 @@ class OpenGLRenderer:
         self._instance_color2_buffer = None
         self._instance_count = 0
         self._wp_instance_ids = None
+        self._wp_instance_custom_ids = None
+        self._np_instance_visible = None
         self._instance_ids = None
         self._inverse_instance_ids = None
         self._wp_instance_transforms = None
@@ -884,6 +990,7 @@ class OpenGLRenderer:
         self._tile_projection_matrices = None
 
         self._frame_texture = None
+        self._frame_depth_texture = None
         self._frame_fbo = None
         self._frame_pbo = None
 
@@ -899,6 +1006,8 @@ class OpenGLRenderer:
 
         gl.glClearColor(*self.background_color, 1)
         gl.glEnable(gl.GL_DEPTH_TEST)
+        gl.glDepthMask(True)
+        gl.glDepthRange(0.0, 1.0)
 
         self._shape_shader = ShaderProgram(
             Shader(shape_vertex_shader, "vertex"), Shader(shape_fragment_shader, "fragment")
@@ -965,20 +1074,22 @@ class OpenGLRenderer:
 
             self._loc_sky_color1 = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("color1"))
             self._loc_sky_color2 = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("color2"))
+            self._loc_sky_far_plane = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("farPlane"))
             gl.glUniform3f(self._loc_sky_color1, *background_color)
             # glUniform3f(self._loc_sky_color2, *np.clip(np.array(background_color)+0.5, 0.0, 1.0))
             gl.glUniform3f(self._loc_sky_color2, 0.8, 0.4, 0.05)
+            gl.glUniform1f(self._loc_sky_far_plane, self.camera_far_plane)
             self._loc_sky_view_pos = gl.glGetUniformLocation(self._sky_shader.id, str_buffer("viewPos"))
             gl.glUniform3f(
                 gl.glGetUniformLocation(self._sky_shader.id, str_buffer("sunDirection")), *self._sun_direction
             )
 
-        # Create VAO, VBO, and EBO
+        # create VAO, VBO, and EBO
         self._sky_vao = gl.GLuint()
         gl.glGenVertexArrays(1, self._sky_vao)
         gl.glBindVertexArray(self._sky_vao)
 
-        vertices, indices = self._create_sphere_mesh(self.camera_far_plane * 0.9, 32, 32)
+        vertices, indices = self._create_sphere_mesh(self.camera_far_plane * 0.9, 32, 32, reverse_winding=True)
         self._sky_tri_count = len(indices)
 
         self._sky_vbo = gl.GLuint()
@@ -991,7 +1102,7 @@ class OpenGLRenderer:
         gl.glBindBuffer(gl.GL_ELEMENT_ARRAY_BUFFER, self._sky_ebo)
         gl.glBufferData(gl.GL_ELEMENT_ARRAY_BUFFER, indices.nbytes, indices.ctypes.data, gl.GL_STATIC_DRAW)
 
-        # Set up vertex attributes
+        # set up vertex attributes
         vertex_stride = vertices.shape[1] * vertices.itemsize
         # positions
         gl.glVertexAttribPointer(0, 3, gl.GL_FLOAT, gl.GL_FALSE, vertex_stride, ctypes.c_void_p(0))
@@ -1025,6 +1136,7 @@ class OpenGLRenderer:
 
         # create frame buffer for rendering to a texture
         self._frame_texture = None
+        self._frame_depth_texture = None
         self._frame_fbo = None
         self._setup_framebuffer()
 
@@ -1072,7 +1184,15 @@ class OpenGLRenderer:
         gl.glUseProgram(self._frame_shader.id)
         self._frame_loc_texture = gl.glGetUniformLocation(self._frame_shader.id, str_buffer("textureSampler"))
 
-        # Unbind the VBO and VAO
+        self._frame_depth_shader = ShaderProgram(
+            Shader(frame_vertex_shader, "vertex"), Shader(frame_depth_fragment_shader, "fragment")
+        )
+        gl.glUseProgram(self._frame_depth_shader.id)
+        self._frame_loc_depth_texture = gl.glGetUniformLocation(
+            self._frame_depth_shader.id, str_buffer("textureSampler")
+        )
+
+        # unbind the VBO and VAO
         gl.glBindBuffer(gl.GL_ARRAY_BUFFER, 0)
         gl.glBindVertexArray(0)
 
@@ -1160,11 +1280,17 @@ class OpenGLRenderer:
         self._instance_color1_buffer = None
         self._instance_color2_buffer = None
         self._wp_instance_ids = None
+        self._wp_instance_custom_ids = None
         self._wp_instance_transforms = None
         self._wp_instance_scalings = None
         self._wp_instance_bodies = None
+        self._np_instance_visible = None
         self._update_shape_instances = False
 
+    def close(self):
+        self.clear()
+        self.window.close()
+
     @property
     def tiled_rendering(self):
         return self._tiled_rendering
@@ -1318,7 +1444,11 @@ class OpenGLRenderer:
         if self._frame_texture is None:
             self._frame_texture = gl.GLuint()
             gl.glGenTextures(1, self._frame_texture)
+        if self._frame_depth_texture is None:
+            self._frame_depth_texture = gl.GLuint()
+            gl.glGenTextures(1, self._frame_depth_texture)
 
+        # set up RGB texture
         gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0)
         gl.glBindBuffer(gl.GL_PIXEL_UNPACK_BUFFER, 0)
         gl.glBindTexture(gl.GL_TEXTURE_2D, self._frame_texture)
@@ -1335,6 +1465,22 @@ class OpenGLRenderer:
         )
         gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR)
         gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR)
+
+        # set up depth texture
+        gl.glBindTexture(gl.GL_TEXTURE_2D, self._frame_depth_texture)
+        gl.glTexImage2D(
+            gl.GL_TEXTURE_2D,
+            0,
+            gl.GL_DEPTH_COMPONENT32,
+            self.screen_width,
+            self.screen_height,
+            0,
+            gl.GL_DEPTH_COMPONENT,
+            gl.GL_FLOAT,
+            None,
+        )
+        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR)
+        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR)
         gl.glBindTexture(gl.GL_TEXTURE_2D, 0)
 
         # create a framebuffer object (FBO)
@@ -1347,15 +1493,9 @@ class OpenGLRenderer:
             gl.glFramebufferTexture2D(
                 gl.GL_FRAMEBUFFER, gl.GL_COLOR_ATTACHMENT0, gl.GL_TEXTURE_2D, self._frame_texture, 0
             )
-
-            self._frame_depth_renderbuffer = gl.GLuint()
-            gl.glGenRenderbuffers(1, self._frame_depth_renderbuffer)
-            gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, self._frame_depth_renderbuffer)
-            gl.glRenderbufferStorage(gl.GL_RENDERBUFFER, gl.GL_DEPTH_COMPONENT, self.screen_width, self.screen_height)
-
-            # attach the depth renderbuffer to the FBO
-            gl.glFramebufferRenderbuffer(
-                gl.GL_FRAMEBUFFER, gl.GL_DEPTH_ATTACHMENT, gl.GL_RENDERBUFFER, self._frame_depth_renderbuffer
+            # attach the depth texture to the FBO as its depth attachment
+            gl.glFramebufferTexture2D(
+                gl.GL_FRAMEBUFFER, gl.GL_DEPTH_ATTACHMENT, gl.GL_TEXTURE_2D, self._frame_depth_texture, 0
             )
 
             if gl.glCheckFramebufferStatus(gl.GL_FRAMEBUFFER) != gl.GL_FRAMEBUFFER_COMPLETE:
@@ -1363,13 +1503,6 @@ class OpenGLRenderer:
                 gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0)
                 sys.exit(1)
 
-            gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, 0)
-        else:
-            # rescale framebuffer
-            gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, self._frame_depth_renderbuffer)
-            gl.glRenderbufferStorage(gl.GL_RENDERBUFFER, gl.GL_DEPTH_COMPONENT, self.screen_width, self.screen_height)
-            gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, 0)
-
         # unbind the FBO (switch back to the default framebuffer)
         gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0)
 
@@ -1379,7 +1512,11 @@ class OpenGLRenderer:
         gl.glBindBuffer(gl.GL_PIXEL_PACK_BUFFER, self._frame_pbo)  # binding to this buffer
 
         # allocate memory for PBO
-        pixels = np.zeros((self.screen_height, self.screen_width, 3), dtype=np.uint8)
+        rgb_bytes_per_pixel = 3
+        depth_bytes_per_pixel = 4
+        pixels = np.zeros(
+            (self.screen_height, self.screen_width, rgb_bytes_per_pixel + depth_bytes_per_pixel), dtype=np.uint8
+        )
         gl.glBufferData(gl.GL_PIXEL_PACK_BUFFER, pixels.nbytes, pixels.ctypes.data, gl.GL_DYNAMIC_DRAW)
         gl.glBindBuffer(gl.GL_PIXEL_PACK_BUFFER, 0)
 
@@ -1431,7 +1568,7 @@ class OpenGLRenderer:
             ))
         elif self._camera_axis == 2:
             self._model_matrix = np.array((
-                self._scaling, 0, 0, 0,
+                -self._scaling, 0, 0, 0,
                 0, 0, self._scaling, 0,
                 0, self._scaling, 0, 0,
                 0, 0, 0, 1
@@ -1502,28 +1639,28 @@ class OpenGLRenderer:
         self._last_time = self.clock_time
         self._frame_speed = update_duration * 100.0
 
-        self.app.event_loop.idle()
-        self.app.platform_event_loop.step(self._frame_dt)
+        # self.app.event_loop.idle()
+        self.app.platform_event_loop.step(self._frame_dt * 1e-3)
 
-        self._skip_frame_counter += 1
-        if self._skip_frame_counter > 100:
-            self._skip_frame_counter = 0
+        if not self.skip_rendering:
+            self._skip_frame_counter += 1
+            if self._skip_frame_counter > 100:
+                self._skip_frame_counter = 0
 
-        if frame_duration > 0.0:
-            if self._fps_update is None:
-                self._fps_update = 1.0 / frame_duration
-            else:
-                update = 1.0 / frame_duration
-                self._fps_update = (1.0 - self._fps_alpha) * self._fps_update + self._fps_alpha * update
-        if update_duration > 0.0:
-            if self._fps_render is None:
-                self._fps_render = 1.0 / update_duration
-            else:
-                update = 1.0 / update_duration
-                self._fps_render = (1.0 - self._fps_alpha) * self._fps_render + self._fps_alpha * update
+            if frame_duration > 0.0:
+                if self._fps_update is None:
+                    self._fps_update = 1.0 / frame_duration
+                else:
+                    update = 1.0 / frame_duration
+                    self._fps_update = (1.0 - self._fps_alpha) * self._fps_update + self._fps_alpha * update
+            if update_duration > 0.0:
+                if self._fps_render is None:
+                    self._fps_render = 1.0 / update_duration
+                else:
+                    update = 1.0 / update_duration
+                    self._fps_render = (1.0 - self._fps_alpha) * self._fps_render + self._fps_alpha * update
 
-        if not self.skip_rendering:
-            self.app.event_loop._redraw_windows(self._frame_dt)
+            self.app.event_loop._redraw_windows(self._frame_dt * 1e-3)
 
     def _draw(self):
         from pyglet import gl
@@ -1531,9 +1668,13 @@ class OpenGLRenderer:
         # catch key hold events
         self._process_inputs()
 
+        if self.enable_backface_culling:
+            gl.glEnable(gl.GL_CULL_FACE)
+        else:
+            gl.glDisable(gl.GL_CULL_FACE)
+
         if self._frame_fbo is not None:
             gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, self._frame_fbo)
-            gl.glBindBuffer(gl.GL_PIXEL_UNPACK_BUFFER, self._frame_fbo)
 
         gl.glClearColor(*self.background_color, 1)
         gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT)
@@ -1554,11 +1695,19 @@ class OpenGLRenderer:
         gl.glUniformMatrix4fv(self._loc_shape_view, 1, gl.GL_FALSE, view_mat_ptr)
         gl.glUniformMatrix4fv(self._loc_shape_projection, 1, gl.GL_FALSE, projection_mat_ptr)
 
+        if self.render_wireframe:
+            gl.glPolygonMode(gl.GL_FRONT_AND_BACK, gl.GL_LINE)
+
         if self._tiled_rendering:
             self._render_scene_tiled()
         else:
             self._render_scene()
 
+        for cb in self.render_3d_callbacks:
+            cb()
+
+        gl.glPolygonMode(gl.GL_FRONT_AND_BACK, gl.GL_FILL)
+
         gl.glBindBuffer(gl.GL_PIXEL_UNPACK_BUFFER, 0)
         gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0)
         gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT)
@@ -1566,15 +1715,26 @@ class OpenGLRenderer:
 
         # render frame buffer texture to screen
         if self._frame_fbo is not None:
-            with self._frame_shader:
-                gl.glActiveTexture(gl.GL_TEXTURE0)
-                gl.glBindTexture(gl.GL_TEXTURE_2D, self._frame_texture)
-                gl.glUniform1i(self._frame_loc_texture, 0)
+            if self.render_depth:
+                with self._frame_depth_shader:
+                    gl.glActiveTexture(gl.GL_TEXTURE0)
+                    gl.glBindTexture(gl.GL_TEXTURE_2D, self._frame_depth_texture)
+                    gl.glUniform1i(self._frame_loc_depth_texture, 0)
+
+                    gl.glBindVertexArray(self._frame_vao)
+                    gl.glDrawElements(gl.GL_TRIANGLES, len(self._frame_indices), gl.GL_UNSIGNED_INT, None)
+                    gl.glBindVertexArray(0)
+                    gl.glBindTexture(gl.GL_TEXTURE_2D, 0)
+            else:
+                with self._frame_shader:
+                    gl.glActiveTexture(gl.GL_TEXTURE0)
+                    gl.glBindTexture(gl.GL_TEXTURE_2D, self._frame_texture)
+                    gl.glUniform1i(self._frame_loc_texture, 0)
 
-                gl.glBindVertexArray(self._frame_vao)
-                gl.glDrawElements(gl.GL_TRIANGLES, len(self._frame_indices), gl.GL_UNSIGNED_INT, None)
-                gl.glBindVertexArray(0)
-                gl.glBindTexture(gl.GL_TEXTURE_2D, 0)
+                    gl.glBindVertexArray(self._frame_vao)
+                    gl.glDrawElements(gl.GL_TRIANGLES, len(self._frame_indices), gl.GL_UNSIGNED_INT, None)
+                    gl.glBindVertexArray(0)
+                    gl.glBindTexture(gl.GL_TEXTURE_2D, 0)
 
         # check for OpenGL errors
         # check_gl_error()
@@ -1597,6 +1757,9 @@ Instances: {len(self._instances)}"""
             self._info_label.y = self.screen_height - 5
             self._info_label.draw()
 
+        for cb in self.render_2d_callbacks:
+            cb()
+
     def _draw_grid(self, is_tiled=False):
         from pyglet import gl
 
@@ -1694,6 +1857,9 @@ Instances: {len(self._instances)}"""
         gl.glBindVertexArray(0)
 
     def _mouse_drag_callback(self, x, y, dx, dy, buttons, modifiers):
+        if not self.enable_mouse_interaction:
+            return
+
         import pyglet
 
         if buttons & pyglet.window.mouse.LEFT:
@@ -1713,6 +1879,9 @@ Instances: {len(self._instances)}"""
             self.update_view_matrix()
 
     def _scroll_callback(self, x, y, scroll_x, scroll_y):
+        if not self.enable_mouse_interaction:
+            return
+
         self.camera_fov -= scroll_y
         self.camera_fov = max(min(self.camera_fov, 90.0), 15.0)
         self.update_projection_matrix()
@@ -1739,8 +1908,11 @@ Instances: {len(self._instances)}"""
     def _key_press_callback(self, symbol, modifiers):
         import pyglet
 
+        if not self.enable_keyboard_interaction:
+            return
+
         if symbol == pyglet.window.key.ESCAPE:
-            self.window.close()
+            self.close()
         if symbol == pyglet.window.key.SPACE:
             self.paused = not self.paused
         if symbol == pyglet.window.key.TAB:
@@ -1751,6 +1923,18 @@ Instances: {len(self._instances)}"""
             self.draw_grid = not self.draw_grid
         if symbol == pyglet.window.key.I:
             self.show_info = not self.show_info
+        if symbol == pyglet.window.key.X:
+            self.render_wireframe = not self.render_wireframe
+        if symbol == pyglet.window.key.T:
+            self.render_depth = not self.render_depth
+        if symbol == pyglet.window.key.B:
+            self.enable_backface_culling = not self.enable_backface_culling
+
+        for cb in self._key_callbacks:
+            cb(symbol, modifiers)
+
+    def register_key_press_callback(self, callback):
+        self._key_callbacks.append(callback)
 
     def _window_resize_callback(self, width, height):
         self._first_mouse = True
@@ -1808,7 +1992,17 @@ Instances: {len(self._instances)}"""
         return shape
 
     def add_shape_instance(
-        self, name: str, shape: int, body, pos, rot, scale=(1.0, 1.0, 1.0), color1=None, color2=None
+        self,
+        name: str,
+        shape: int,
+        body,
+        pos,
+        rot,
+        scale=(1.0, 1.0, 1.0),
+        color1=None,
+        color2=None,
+        custom_index: int = -1,
+        visible: bool = True,
     ):
         if color1 is None:
             color1 = self._shapes[shape][2]
@@ -1817,8 +2011,9 @@ Instances: {len(self._instances)}"""
         instance = len(self._instances)
         self._shape_instances[shape].append(instance)
         body = self._resolve_body_id(body)
-        self._instances[name] = (instance, body, shape, [*pos, *rot], scale, color1, color2)
+        self._instances[name] = (instance, body, shape, [*pos, *rot], scale, color1, color2, visible)
         self._instance_shape[instance] = shape
+        self._instance_custom_ids[instance] = custom_index
         self._add_shape_instances = True
         self._instance_count = len(self._instances)
         return instance
@@ -1843,7 +2038,7 @@ Instances: {len(self._instances)}"""
             gl.glDeleteBuffers(1, self._instance_color1_buffer)
             gl.glDeleteBuffers(1, self._instance_color2_buffer)
 
-        # Create instance buffer and bind it as an instanced array
+        # create instance buffer and bind it as an instanced array
         self._instance_transform_gl_buffer = gl.GLuint()
         gl.glGenBuffers(1, self._instance_transform_gl_buffer)
         gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self._instance_transform_gl_buffer)
@@ -1851,7 +2046,7 @@ Instances: {len(self._instances)}"""
         transforms = np.tile(np.diag(np.ones(4, dtype=np.float32)), (len(self._instances), 1, 1))
         gl.glBufferData(gl.GL_ARRAY_BUFFER, transforms.nbytes, transforms.ctypes.data, gl.GL_DYNAMIC_DRAW)
 
-        # Create CUDA buffer for instance transforms
+        # create CUDA buffer for instance transforms
         self._instance_transform_cuda_buffer = wp.RegisteredGLBuffer(
             int(self._instance_transform_gl_buffer.value), self._device
         )
@@ -1879,10 +2074,13 @@ Instances: {len(self._instances)}"""
         gl.glBindBuffer(gl.GL_ARRAY_BUFFER, self._instance_color2_buffer)
         gl.glBufferData(gl.GL_ARRAY_BUFFER, colors2.nbytes, colors2.ctypes.data, gl.GL_STATIC_DRAW)
 
-        # Set up instance attribute pointers
+        # set up instance attribute pointers
         matrix_size = transforms[0].nbytes
 
         instance_ids = []
+        instance_custom_ids = []
+        instance_visible = []
+        instances = list(self._instances.values())
         inverse_instance_ids = {}
         instance_count = 0
         for shape, (vao, vbo, ebo, tri_count, vertex_cuda_buffer) in self._shape_gl_buffers.items():
@@ -1912,27 +2110,45 @@ Instances: {len(self._instances)}"""
             for i in self._shape_instances[shape]:
                 inverse_instance_ids[i] = instance_count
                 instance_count += 1
+                instance_custom_ids.append(self._instance_custom_ids[i])
+                instance_visible.append(instances[i][7])
 
         # trigger update to the instance transforms
         self._update_shape_instances = True
 
         self._wp_instance_ids = wp.array(instance_ids, dtype=wp.int32, device=self._device)
+        self._wp_instance_custom_ids = wp.array(instance_custom_ids, dtype=wp.int32, device=self._device)
+        self._np_instance_visible = np.array(instance_visible)
         self._instance_ids = instance_ids
         self._inverse_instance_ids = inverse_instance_ids
 
         gl.glBindVertexArray(0)
 
-    def update_shape_instance(self, name, pos, rot, color1=None, color2=None):
+    def update_shape_instance(self, name, pos, rot, color1=None, color2=None, visible=None):
         """Update the instance transform of the shape
 
         Args:
             name: The name of the shape
             pos: The position of the shape
             rot: The rotation of the shape
+            color1: The first color of the checker pattern
+            color2: The second color of the checker pattern
+            visible: Whether the shape is visible
         """
         if name in self._instances:
-            i, body, shape, _, scale, old_color1, old_color2 = self._instances[name]
-            self._instances[name] = (i, body, shape, [*pos, *rot], scale, color1 or old_color1, color2 or old_color2)
+            i, body, shape, _, scale, old_color1, old_color2, v = self._instances[name]
+            if visible is None:
+                visible = v
+            self._instances[name] = (
+                i,
+                body,
+                shape,
+                [*pos, *rot],
+                scale,
+                color1 or old_color1,
+                color2 or old_color2,
+                visible,
+            )
             self._update_shape_instances = True
             return True
         return False
@@ -2000,9 +2216,27 @@ Instances: {len(self._instances)}"""
             self.clear()
             self.app.event_loop.exit()
 
-    def get_pixels(self, target_image: wp.array, split_up_tiles=True):
+    def get_pixels(self, target_image: wp.array, split_up_tiles=True, mode="rgb"):
+        """
+        Read the pixels from the frame buffer (RGB or depth are supported) into the given array.
+
+        If `split_up_tiles` is False, array must be of shape (screen_height, screen_width, 3) for RGB mode or
+        (screen_height, screen_width, 1) for depth mode.
+        If `split_up_tiles` is True, the pixels will be split up into tiles (see :attr:`tile_width` and :attr:`tile_height` for dimensions):
+        array must be of shape (num_tiles, tile_height, tile_width, 3) for RGB mode or (num_tiles, tile_height, tile_width, 1) for depth mode.
+
+        Args:
+            target_image (array): The array to read the pixels into. Must have float32 as dtype and be on a CUDA device.
+            split_up_tiles (bool): Whether to split up the viewport into tiles, see :meth:`setup_tiled_rendering`.
+            mode (str): can be either "rgb" or "depth"
+
+        Returns:
+            bool: Whether the pixels were successfully read.
+        """
         from pyglet import gl
 
+        channels = 3 if mode == "rgb" else 1
+
         if split_up_tiles:
             assert (
                 self._tile_width is not None and self._tile_height is not None
@@ -2017,20 +2251,26 @@ Instances: {len(self._instances)}"""
                 self.num_tiles,
                 self._tile_height,
                 self._tile_width,
-                3,
-            ), f"Shape of `target_image` array does not match {self.num_tiles} x {self.screen_height} x {self.screen_width} x 3"
+                channels,
+            ), f"Shape of `target_image` array does not match {self.num_tiles} x {self.screen_height} x {self.screen_width} x {channels}"
         else:
             assert target_image.shape == (
                 self.screen_height,
                 self.screen_width,
-                3,
-            ), f"Shape of `target_image` array does not match {self.screen_height} x {self.screen_width} x 3"
+                channels,
+            ), f"Shape of `target_image` array does not match {self.screen_height} x {self.screen_width} x {channels}"
 
         gl.glBindBuffer(gl.GL_PIXEL_PACK_BUFFER, self._frame_pbo)
-        gl.glBindTexture(gl.GL_TEXTURE_2D, self._frame_texture)
+        if mode == "rgb":
+            gl.glBindTexture(gl.GL_TEXTURE_2D, self._frame_texture)
+        if mode == "depth":
+            gl.glBindTexture(gl.GL_TEXTURE_2D, self._frame_depth_texture)
         try:
             # read screen texture into PBO
-            gl.glGetTexImage(gl.GL_TEXTURE_2D, 0, gl.GL_RGB, gl.GL_UNSIGNED_BYTE, ctypes.c_void_p(0))
+            if mode == "rgb":
+                gl.glGetTexImage(gl.GL_TEXTURE_2D, 0, gl.GL_RGB, gl.GL_UNSIGNED_BYTE, ctypes.c_void_p(0))
+            elif mode == "depth":
+                gl.glGetTexImage(gl.GL_TEXTURE_2D, 0, gl.GL_DEPTH_COMPONENT, gl.GL_FLOAT, ctypes.c_void_p(0))
         except gl.GLException:
             # this can happen if the window is closed/being moved to a different display
             gl.glBindTexture(gl.GL_TEXTURE_2D, 0)
@@ -2043,63 +2283,54 @@ Instances: {len(self._instances)}"""
             int(self._frame_pbo.value), self._device, wp.RegisteredGLBuffer.WRITE_DISCARD
         )
         screen_size = self.screen_height * self.screen_width
-        img = pbo_buffer.map(dtype=wp.uint8, shape=(screen_size * 3))
+        if mode == "rgb":
+            img = pbo_buffer.map(dtype=wp.uint8, shape=(screen_size * channels))
+        elif mode == "depth":
+            img = pbo_buffer.map(dtype=wp.float32, shape=(screen_size * channels))
         img = img.to(target_image.device)
         if split_up_tiles:
             positions = wp.array(self._tile_viewports, ndim=2, dtype=wp.int32, device=target_image.device)
-            wp.launch(
-                copy_frame_tiles,
-                dim=(self.num_tiles, self._tile_width, self._tile_height),
-                inputs=[img, positions, self.screen_width, self.screen_height, self._tile_height],
-                outputs=[target_image],
-                device=target_image.device,
-            )
+            if mode == "rgb":
+                wp.launch(
+                    copy_rgb_frame_tiles,
+                    dim=(self.num_tiles, self._tile_width, self._tile_height),
+                    inputs=[img, positions, self.screen_width, self.screen_height, self._tile_height],
+                    outputs=[target_image],
+                    device=target_image.device,
+                )
+            elif mode == "depth":
+                wp.launch(
+                    copy_depth_frame_tiles,
+                    dim=(self.num_tiles, self._tile_width, self._tile_height),
+                    inputs=[
+                        img,
+                        positions,
+                        self.screen_width,
+                        self.screen_height,
+                        self._tile_height,
+                        self.camera_near_plane,
+                        self.camera_far_plane,
+                    ],
+                    outputs=[target_image],
+                    device=target_image.device,
+                )
         else:
-            wp.launch(
-                copy_frame,
-                dim=(self.screen_width, self.screen_height),
-                inputs=[img, self.screen_width, self.screen_height],
-                outputs=[target_image],
-                device=target_image.device,
-            )
-        pbo_buffer.unmap()
-        return True
-
-    def get_tile_pixels(self, tile_id: int, target_image: wp.array):
-        from pyglet import gl
-
-        viewport = self._tile_viewports[tile_id]
-        assert target_image.shape == (
-            viewport[3],
-            viewport[2],
-            3,
-        ), f"Shape of `target_image` array does not match {viewport[3]} x {viewport[2]} x 3"
-        gl.glBindBuffer(gl.GL_PIXEL_PACK_BUFFER, self._frame_pbo)
-        gl.glBindTexture(gl.GL_TEXTURE_2D, self._frame_texture)
-        try:
-            # read screen texture into PBO
-            gl.glGetTexImage(gl.GL_TEXTURE_2D, 0, gl.GL_RGB, gl.GL_UNSIGNED_BYTE, ctypes.c_void_p(0))
-        except gl.GLException:
-            # this can happen if the window is closed/being moved to a different display
-            gl.glBindTexture(gl.GL_TEXTURE_2D, 0)
-            gl.glBindBuffer(gl.GL_PIXEL_PACK_BUFFER, 0)
-            return False
-        gl.glBindTexture(gl.GL_TEXTURE_2D, 0)
-        gl.glBindBuffer(gl.GL_PIXEL_PACK_BUFFER, 0)
-
-        pbo_buffer = wp.RegisteredGLBuffer(
-            int(self._frame_pbo.value), self._device, wp.RegisteredGLBuffer.WRITE_DISCARD
-        )
-        screen_size = self.screen_height * self.screen_width
-        img = pbo_buffer.map(dtype=wp.uint8, shape=(screen_size * 3))
-        img = img.to(target_image.device)
-        wp.launch(
-            copy_frame_tiles,
-            dim=(self.num_tiles, self._tile_width, self._tile_height),
-            inputs=[img, viewport[0], viewport[1], self.screen_width, self.screen_height, self._tile_height],
-            outputs=[target_image],
-            device=target_image.device,
-        )
+            if mode == "rgb":
+                wp.launch(
+                    copy_rgb_frame,
+                    dim=(self.screen_width, self.screen_height),
+                    inputs=[img, self.screen_width, self.screen_height],
+                    outputs=[target_image],
+                    device=target_image.device,
+                )
+            elif mode == "depth":
+                wp.launch(
+                    copy_depth_frame,
+                    dim=(self.screen_width, self.screen_height),
+                    inputs=[img, self.screen_width, self.screen_height, self.camera_near_plane, self.camera_far_plane],
+                    outputs=[target_image],
+                    device=target_image.device,
+                )
         pbo_buffer.unmap()
         return True
 
@@ -2515,6 +2746,7 @@ Instances: {len(self._instances)}"""
             wp_points = wp.array(points, dtype=wp.vec3, device=self._device)
 
         if name not in self._shape_instancers:
+            np_points = points.numpy() if isinstance(points, wp.array) else points
             instancer = ShapeInstancer(self._shape_shader, self._device)
             radius_is_scalar = np.isscalar(radius)
             if radius_is_scalar:
@@ -2527,12 +2759,13 @@ Instances: {len(self._instances)}"""
                 color = colors[0]
             instancer.register_shape(vertices, indices, color, color)
             scalings = None if radius_is_scalar else np.tile(radius, (3, 1)).T
-            instancer.allocate_instances(np.array(points), colors1=colors, colors2=colors, scalings=scalings)
+            instancer.allocate_instances(np_points, colors1=colors, colors2=colors, scalings=scalings)
             self._shape_instancers[name] = instancer
         else:
             instancer = self._shape_instancers[name]
             if len(points) != instancer.num_instances:
-                instancer.allocate_instances(np.array(points))
+                np_points = points.numpy() if isinstance(points, wp.array) else points
+                instancer.allocate_instances(np_points)
 
         with instancer:
             wp.launch(
@@ -2550,7 +2783,7 @@ Instances: {len(self._instances)}"""
         if name not in self._shape_instancers:
             instancer = ShapeInstancer(self._shape_shader, self._device)
             vertices, indices = self._create_capsule_mesh(radius, 0.5)
-            if color is None:
+            if color is None or isinstance(color, list) and len(color) > 0 and isinstance(color[0], list):
                 color = tab10_color_map(len(self._shape_geo_hash))
             instancer.register_shape(vertices, indices, color, color)
             instancer.allocate_instances(np.zeros((len(lines), 3)))
@@ -2570,7 +2803,7 @@ Instances: {len(self._instances)}"""
                 device=self._device,
             )
 
-    def render_line_list(self, name, vertices, indices, color, radius):
+    def render_line_list(self, name: str, vertices, indices, color: tuple = None, radius: float = 0.01):
         """Add a line list as a set of capsules
 
         Args:
@@ -2585,7 +2818,7 @@ Instances: {len(self._instances)}"""
         lines = np.array(lines)
         self._render_lines(name, lines, color, radius)
 
-    def render_line_strip(self, name: str, vertices, color: tuple, radius: float = 0.01):
+    def render_line_strip(self, name: str, vertices, color: tuple = None, radius: float = 0.01):
         """Add a line strip as a set of capsules
 
         Args:
@@ -2620,7 +2853,12 @@ Instances: {len(self._instances)}"""
         cuda_buffer.unmap()
 
     @staticmethod
-    def _create_sphere_mesh(radius=1.0, num_latitudes=default_num_segments, num_longitudes=default_num_segments):
+    def _create_sphere_mesh(
+        radius=1.0,
+        num_latitudes=default_num_segments,
+        num_longitudes=default_num_segments,
+        reverse_winding=False,
+    ):
         vertices = []
         indices = []
 
@@ -2648,7 +2886,10 @@ Instances: {len(self._instances)}"""
                 first = i * (num_longitudes + 1) + j
                 second = first + num_longitudes + 1
 
-                indices.extend([first, second, first + 1, second, second + 1, first + 1])
+                if reverse_winding:
+                    indices.extend([first, second, first + 1, second, second + 1, first + 1])
+                else:
+                    indices.extend([first, first + 1, second, second, first + 1, second + 1])
 
         return np.array(vertices, dtype=np.float32), np.array(indices, dtype=np.uint32)
 
@@ -2731,7 +2972,7 @@ Instances: {len(self._instances)}"""
             top_radius = radius
         side_slope = -np.arctan2(top_radius - radius, 2 * half_height)
 
-        # Create the cylinder base and top vertices
+        # create the cylinder base and top vertices
         for j in (-1, 1):
             center_index = max(j, 0)
             if j == 1:
@@ -2759,11 +3000,10 @@ Instances: {len(self._instances)}"""
                 vertex = np.hstack([position[[x_dir, y_dir, z_dir]], normal[[x_dir, y_dir, z_dir]], uv])
                 cap_vertices.append(vertex)
 
-                indices.extend(
-                    [center_index, i + center_index * segments + 2, (i + 1) % segments + center_index * segments + 2]
-                )
+                cs = center_index * segments
+                indices.extend([center_index, i + cs + 2, (i + 1) % segments + cs + 2][::-j])
 
-        # Create the cylinder side indices
+        # create the cylinder side indices
         for i in range(segments):
             index1 = len(cap_vertices) + i + segments
             index2 = len(cap_vertices) + ((i + 1) % segments) + segments
@@ -2841,16 +3081,16 @@ Instances: {len(self._instances)}"""
         indices = [
             0, 1, 2,
             0, 2, 3,
-            4, 5, 6,
-            4, 6, 7,
-            8, 9, 10,
-            8, 10, 11,
+            4, 6, 5,
+            4, 7, 6,
+            8, 10, 9,
+            8, 11, 10,
             12, 13, 14,
             12, 14, 15,
             16, 17, 18,
             16, 18, 19,
-            20, 21, 22,
-            20, 22, 23,
+            20, 22, 21,
+            20, 23, 22,
         ]
         # fmt: on
         return np.array(vertices, dtype=np.float32), np.array(indices, dtype=np.uint32)