Merge branch 'main' into diegoferigo/fix_mjcf_with_texture

Author: eric-heiden

newton/_src/viewer/viewer_gl.py

@@ -15,10 +15,12 @@
 
 from __future__ import annotations
 
+import ctypes
 import time
 
 import numpy as np
 import warp as wp
+import warp.render.render_opengl
 
 import newton as nt
 from newton.selection import ArticulationView
@@ -127,6 +129,10 @@ def __init__(self, width=1920, height=1080, vsync=False, headless=False):
         # positions: "side", "stats", "free"
         self._ui_callbacks = {"side": [], "stats": [], "free": []}
 
+        # Initialize PBO (Pixel Buffer Object) resources used in the `get_frame` method.
+        self._pbo = None
+        self._wp_pbo = None
+
         self.set_model(None)
 
     def register_ui_callback(self, callback, position="side"):
@@ -475,6 +481,83 @@ def _update(self):
 
         self.renderer.present()
 
+    def get_frame(self, target_image: wp.array | None = None) -> wp.array:
+        """
+        Retrieve the last rendered frame.
+
+        This method uses OpenGL Pixel Buffer Objects (PBO) and CUDA interoperability
+        to transfer pixel data entirely on the GPU, avoiding expensive CPU-GPU transfers.
+
+        Args:
+            target_image (wp.array, optional):
+                Optional pre-allocated Warp array with shape `(height, width, 3)`
+                and dtype `wp.uint8`. If `None`, a new array will be created.
+
+        Returns:
+            wp.array: GPU array containing RGB image data with shape `(height, width, 3)`
+                and dtype `wp.uint8`. Origin is top-left (OpenGL's bottom-left is flipped).
+        """
+
+        gl = RendererGL.gl
+        w, h = self.renderer._screen_width, self.renderer._screen_height
+
+        # Lazy initialization of PBO (Pixel Buffer Object).
+        if self._pbo is None:
+            pbo_id = (gl.GLuint * 1)()
+            gl.glGenBuffers(1, pbo_id)
+            self._pbo = pbo_id[0]
+
+            # Allocate PBO storage.
+            gl.glBindBuffer(gl.GL_PIXEL_PACK_BUFFER, self._pbo)
+            gl.glBufferData(gl.GL_PIXEL_PACK_BUFFER, gl.GLsizeiptr(w * h * 3), None, gl.GL_STREAM_READ)
+            gl.glBindBuffer(gl.GL_PIXEL_PACK_BUFFER, 0)
+
+            # Register with CUDA.
+            self._wp_pbo = wp.RegisteredGLBuffer(
+                gl_buffer_id=int(self._pbo),
+                device=self.device,
+                flags=wp.RegisteredGLBuffer.READ_ONLY,
+            )
+
+            # Set alignment once.
+            gl.glPixelStorei(gl.GL_PACK_ALIGNMENT, 1)
+
+        # GPU-to-GPU readback into PBO.
+        assert self.renderer._frame_fbo is not None
+        gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, self.renderer._frame_fbo)
+        gl.glBindBuffer(gl.GL_PIXEL_PACK_BUFFER, self._pbo)
+        gl.glReadPixels(0, 0, w, h, gl.GL_RGB, gl.GL_UNSIGNED_BYTE, ctypes.c_void_p(0))
+        gl.glBindBuffer(gl.GL_PIXEL_PACK_BUFFER, 0)
+        gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0)
+
+        # Map PBO buffer and copy using RGB kernel.
+        assert self._wp_pbo is not None
+        buf = self._wp_pbo.map(dtype=wp.uint8, shape=(w * h * 3,))
+
+        if target_image is None:
+            target_image = wp.empty(
+                shape=(h, w, 3),
+                dtype=wp.uint8,  # pyright: ignore[reportArgumentType]
+                device=self.device,
+            )
+
+        if target_image.shape != (h, w, 3):
+            raise ValueError(f"Shape of `target_image` must be ({h}, {w}, 3), got {target_image.shape}")
+
+        # Launch the RGB kernel.
+        wp.launch(
+            warp.render.render_opengl.copy_rgb_frame_uint8,
+            dim=(w, h),
+            inputs=[buf, w, h],
+            outputs=[target_image],
+            device=self.device,
+        )
+
+        # Unmap the PBO buffer.
+        self._wp_pbo.unmap()
+
+        return target_image
+
     @override
     def is_running(self) -> bool:
         """