Add support for storing last rendered frame in ViewerGL

Author: diegoferigo-rai

newton/_src/viewer/gl/opengl.py

@@ -19,6 +19,7 @@
 import sys
 
 import numpy as np
+import numpy.typing as npt
 import warp as wp
 
 from newton.utils import create_sphere_mesh
@@ -922,7 +923,7 @@ def update(self):
             self.app.platform_event_loop.step(0.001)  # 1ms app polling latency
             self.window.dispatch_events()
 
-    def render(self, camera, objects, lines=None):
+    def render(self, camera, objects, lines=None, return_frame: bool = False) -> npt.NDArray[np.uint8] | None:
         gl = RendererGL.gl
         self._make_current()
 
@@ -991,6 +992,21 @@ def render(self, camera, objects, lines=None):
                 gl.GL_NEAREST,
             )
 
+        # -----------------------------------
+        # Read back frame buffer if requested
+        # -----------------------------------
+
+        frame: npt.NDArray[np.uint8] | None = None
+
+        if return_frame:
+            gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, self._frame_fbo)
+            buffer = (gl.GLubyte * (4 * self._screen_width * self._screen_height))()
+            gl.glReadPixels(0, 0, self._screen_width, self._screen_height, gl.GL_RGBA, gl.GL_UNSIGNED_BYTE, buffer)
+            gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0)
+
+            frame = np.frombuffer(buffer, dtype=np.uint8).reshape((self._screen_height, self._screen_width, 4))
+            frame = np.flipud(frame)
+
         # ------------------------------------------------------------------
         # Draw resolved texture to the screen
         # ------------------------------------------------------------------
@@ -1018,6 +1034,8 @@ def render(self, camera, objects, lines=None):
         err = gl.glGetError()
         assert err == gl.GL_NO_ERROR, hex(err)
 
+        return frame
+
     def present(self):
         if not self.headless:
             self.window.flip()

newton/_src/viewer/viewer_gl.py

@@ -18,6 +18,7 @@
 import time
 
 import numpy as np
+import numpy.typing as npt
 import warp as wp
 
 import newton as nt
@@ -53,7 +54,7 @@ class ViewerGL(ViewerBase):
         - Extensible logging of meshes, lines, points, and arrays for custom visualization.
     """
 
-    def __init__(self, width=1920, height=1080, vsync=False, headless=False):
+    def __init__(self, width=1920, height=1080, vsync=False, headless=False, save_frame: bool = False):
         """
         Initialize the OpenGL viewer and UI.
 
@@ -62,6 +63,7 @@ def __init__(self, width=1920, height=1080, vsync=False, headless=False):
             height (int): Window height in pixels.
             vsync (bool): Enable vertical sync.
             headless (bool): Run in headless mode (no window).
+            save_frame (bool): If True, store internally the last rendered frame.
         """
         super().__init__()
 
@@ -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": []}
 
+        # Buffer to store the last rendered frame
+        self._save_frame = save_frame
+        self._frame = np.empty(0, dtype=np.uint8)
+
         self.set_model(None)
 
     def register_ui_callback(self, callback, position="side"):
@@ -425,7 +431,7 @@ def end_frame(self):
         destroyed results in a crash (access violation).  Therefore we check
         whether an exit was requested and early-out before touching GL if so.
         """
-        self._update()
+        return self._update()
 
     @override
     def apply_forces(self, state):
@@ -459,7 +465,10 @@ def _update(self):
             return
 
         # Render the scene and present it
-        self.renderer.render(self.camera, self.objects, self.lines)
+        frame = self.renderer.render(self.camera, self.objects, self.lines, return_frame=self._save_frame)
+
+        if self._save_frame:
+            self._frame = frame
 
         # Always update FPS tracking, even if UI is hidden
         self._update_fps()
@@ -475,6 +484,19 @@ def _update(self):
 
         self.renderer.present()
 
+    def get_frame(self) -> npt.NDArray[np.uint8]:
+        """
+        Get the last rendered frame as a numpy array.
+
+        Returns:
+            np.ndarray: The last rendered frame in RGBA format.
+        """
+
+        if not self._save_frame:
+            raise RuntimeError("Frame saving is not enabled. Initialize with `save_frame=True`.")
+
+        return self._frame
+
     @override
     def is_running(self) -> bool:
         """