Added option to render a texture (e.g. Sensor output) as main image in the viewer

Author: daniela-hase

changelog/+viewer-main-image-18cc9933.added.md

@@ -0,0 +1,3 @@
+Add `ViewerGL.log_main_image()` to display a logged image (e.g. a sensor output
+texture) as the main viewer surface for a frame, skipping the 3D scene render
+while keeping the UI available.

docs/guide/visualization.rst

@@ -187,6 +187,19 @@ Warp array on the viewer device:
     # Returns a wp.array with shape (height, width, 3), dtype wp.uint8
     frame = viewer.get_frame()
 
+**Main image display:**
+
+:meth:`~newton.viewer.ViewerGL.log_main_image` logs an image and displays it as
+the main viewer surface for the current frame. While a main image is logged,
+:class:`~newton.viewer.ViewerGL` skips the normal 3D scene render and draws that
+image directly to the window, with the UI still available on top. If a later
+frame does not call ``log_main_image()``, the viewer returns to normal 3D scene
+rendering for that frame:
+
+.. code-block:: python
+
+    viewer.log_main_image("sensor", sensor_rgba)
+
 **Custom UI panels:**
 
 :meth:`~newton.viewer.ViewerGL.register_ui_callback` adds custom imgui UI elements to the viewer.

newton/_src/viewer/gl/image_logger.py

@@ -326,6 +326,19 @@ class LoggedImage:
     window_initialized: bool = False
 
 
+@dataclass(frozen=True)
+class LoggedImageTexture:
+    """Texture metadata needed to draw a logged image atlas."""
+
+    texture_id: int
+    texture_width: int
+    texture_height: int
+    tile_count: int
+    tile_width: int
+    tile_height: int
+    atlas_cols: int
+
+
 class ImageLogger:
     """Owns GL resources for images logged via :meth:`~newton.viewer.ViewerBase.log_image`.
 
@@ -408,16 +421,20 @@ def log(self, name: str, image: wp.array[Any] | np.ndarray) -> None:
         entry.atlas_cols, entry.atlas_rows = atlas_cols, atlas_rows
         entry.tile_aspect = h / w
 
-    def draw(self) -> None:
+    def draw(self, *, hidden_name: str | None = None) -> None:
         """Draw the selected image window (if any).
 
         Called once per frame inside the viewer's ImGui frame block.
         At most one window is visible at a time; selection is driven by
         :meth:`draw_controls`.
+
+        Args:
+            hidden_name: Optional selected image name to suppress, used when
+                that image is already drawn as the main viewer surface.
         """
         from imgui_bundle import imgui
 
-        if self._selected is None:
+        if self._selected is None or self._selected == hidden_name:
             return
         entry = self._images.get(self._selected)
         if entry is None or entry.n == 0 or entry.tex_id == 0:
@@ -524,6 +541,29 @@ def draw_controls(self) -> None:
         if changed:
             self._selected = None if new_idx == 0 else names[new_idx - 1]
 
+    def get_texture(self, name: str) -> LoggedImageTexture | None:
+        """Return texture metadata for a logged image.
+
+        Args:
+            name: Image name previously passed to :meth:`log`.
+
+        Returns:
+            Texture metadata for the packed image atlas, or ``None`` when the
+            image has not been logged or has no live GL texture yet.
+        """
+        entry = self._images.get(name)
+        if entry is None or entry.tex_id == 0 or entry.tex_w <= 0 or entry.tex_h <= 0:
+            return None
+        return LoggedImageTexture(
+            texture_id=entry.tex_id,
+            texture_width=entry.tex_w,
+            texture_height=entry.tex_h,
+            tile_count=entry.n,
+            tile_width=entry.w,
+            tile_height=entry.h,
+            atlas_cols=entry.atlas_cols,
+        )
+
     def clear(self) -> None:
         """Destroy all GL resources. Idempotent."""
         for entry in list(self._images.values()):

newton/_src/viewer/gl/opengl.py

@@ -53,6 +53,27 @@ def check_gl_error():
         print(f"Called from: {''.join(stack[-2:-1])}")
 
 
+def _texture_tile_uv_rect(
+    tile_index: int,
+    tile_width: int,
+    tile_height: int,
+    texture_width: int,
+    texture_height: int,
+    atlas_cols: int,
+) -> tuple[float, float, float, float]:
+    """Return a V-flipped UV rect for a row-major image atlas tile."""
+    atlas_cols = max(1, int(atlas_cols))
+    atlas_row, atlas_col = divmod(tile_index, atlas_cols)
+    u_step = tile_width / float(texture_width)
+    v_step = tile_height / float(texture_height)
+    return (
+        atlas_col * u_step,
+        (atlas_row + 1) * v_step,
+        (atlas_col + 1) * u_step,
+        atlas_row * v_step,
+    )
+
+
 def _upload_texture_from_file(gl, texture_image: np.ndarray) -> int:
     image = normalize_texture(
         texture_image,
@@ -1334,6 +1355,108 @@ def render(self, camera, objects, lines=None, wireframe_shapes=None, arrows=None
         err = gl.glGetError()
         assert err == gl.GL_NO_ERROR, hex(err)
 
+    def render_texture(
+        self,
+        texture_id: int | None,
+        texture_width: int,
+        texture_height: int,
+        *,
+        tile_count: int = 1,
+        tile_width: int | None = None,
+        tile_height: int | None = None,
+        atlas_cols: int = 1,
+        spacing_px: float = 2.0,
+        clear_color: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0),
+    ):
+        """Draw a texture directly to the window without rendering the 3D scene.
+
+        Args:
+            texture_id: OpenGL texture id to draw, or ``None`` to only clear.
+            texture_width: Source texture width in pixels.
+            texture_height: Source texture height in pixels.
+            tile_count: Number of source tiles packed into the texture atlas.
+            tile_width: Width of each atlas tile in pixels. Defaults to
+                ``texture_width`` for single-image textures.
+            tile_height: Height of each atlas tile in pixels. Defaults to
+                ``texture_height`` for single-image textures.
+            atlas_cols: Number of atlas columns used to pack the source tiles.
+            spacing_px: Spacing between displayed tiles in pixels.
+            clear_color: Window clear color.
+        """
+        gl = RendererGL.gl
+        self._make_current()
+
+        screen_w = max(int(self._screen_width), 1)
+        screen_h = max(int(self._screen_height), 1)
+
+        gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0)
+        gl.glClearColor(*clear_color)
+        gl.glDisable(gl.GL_DEPTH_TEST)
+        gl.glDepthMask(True)
+        gl.glDisable(gl.GL_BLEND)
+        gl.glViewport(0, 0, screen_w, screen_h)
+        gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT)
+        gl.glDepthMask(False)
+
+        if texture_id is None or texture_id == 0 or texture_width <= 0 or texture_height <= 0:
+            gl.glDepthMask(True)
+            return
+
+        from .image_logger import compute_grid_layout  # noqa: PLC0415
+
+        tile_count = max(1, int(tile_count))
+        tile_width = int(tile_width or texture_width)
+        tile_height = int(tile_height or texture_height)
+        atlas_cols = max(1, int(atlas_cols))
+        spacing_px = max(0.0, float(spacing_px if tile_count > 1 else 0.0))
+
+        rows, cols, cell_w, cell_h = compute_grid_layout(
+            tile_count,
+            tile_height / float(max(tile_width, 1)),
+            float(screen_w),
+            float(screen_h),
+            spacing_x=spacing_px,
+            spacing_y=spacing_px,
+        )
+        grid_w = cols * cell_w + max(0, cols - 1) * spacing_px
+        grid_h = rows * cell_h + max(0, rows - 1) * spacing_px
+        origin_x = (screen_w - grid_w) * 0.5
+        origin_y = (screen_h - grid_h) * 0.5
+
+        gl.glActiveTexture(gl.GL_TEXTURE0)
+        gl.glBindTexture(gl.GL_TEXTURE_2D, int(texture_id))
+        with self._frame_shader:
+            for i in range(tile_count):
+                display_row, display_col = divmod(i, cols)
+                draw_x = max(0, int(round(origin_x + display_col * (cell_w + spacing_px))))
+                draw_y = max(
+                    0,
+                    int(round(origin_y + (rows - 1 - display_row) * (cell_h + spacing_px))),
+                )
+                draw_w = max(1, int(round(cell_w)))
+                draw_h = max(1, int(round(cell_h)))
+                uv_rect = _texture_tile_uv_rect(
+                    i,
+                    tile_width,
+                    tile_height,
+                    texture_width,
+                    texture_height,
+                    atlas_cols,
+                )
+
+                gl.glViewport(draw_x, draw_y, draw_w, draw_h)
+                self._frame_shader.update(0, uv_rect)
+                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.glViewport(0, 0, screen_w, screen_h)
+        gl.glDepthMask(True)
+
+        err = gl.glGetError()
+        assert err == gl.GL_NO_ERROR, hex(err)
+
     def present(self):
         if not self.headless:
             if self._dwm_flush is not None and self.window._interval:

newton/_src/viewer/gl/shaders.py

@@ -461,9 +461,11 @@
 out vec4 FragColor;
 
 uniform sampler2D texture_sampler;
+uniform vec4 uv_rect;
 
 void main() {
-    FragColor = texture(texture_sampler, TexCoord);
+    vec2 uv = mix(uv_rect.xy, uv_rect.zw, TexCoord);
+    FragColor = texture(texture_sampler, uv);
 }
 """
 
@@ -696,11 +698,13 @@ def __init__(self, gl):
         # Get uniform locations
         with self:
             self.loc_texture = self._get_uniform_location("texture_sampler")
+            self.loc_uv_rect = self._get_uniform_location("uv_rect")
 
-    def update(self, texture_unit: int = 0):
-        """Update texture uniform."""
+    def update(self, texture_unit: int = 0, uv_rect: tuple[float, float, float, float] = (0.0, 0.0, 1.0, 1.0)):
+        """Update the texture unit and UV sub-rect uniforms."""
         with self:
             self._gl.glUniform1i(self.loc_texture, texture_unit)
+            self._gl.glUniform4f(self.loc_uv_rect, *uv_rect)
 
 
 wireframe_vertex_shader = """

newton/_src/viewer/viewer_gl.py

@@ -246,6 +246,7 @@ def __init__(
             sidebar_width_px=self._sidebar_width_fb_px(),
             dpi_scale=self._dpi_scale(),
         )
+        self._main_image_name: str | None = None
 
         fb_w, fb_h = self.renderer.window.get_framebuffer_size()
         self.camera = Camera(width=fb_w, height=fb_h, up_axis="Z")
@@ -280,7 +281,9 @@ def __init__(
             # Register GL-specific rendering options (sky, shadows, wireframe, colors)
             self.gui.register_ui_callback(self._ui_populate_rendering_panel, position="rendering")
             # Draw image-logger floating windows outside the sidebar window.
-            self.gui.register_ui_callback(lambda _imgui: self._image_logger.draw(), position="free")
+            self.gui.register_ui_callback(
+                lambda _imgui: self._image_logger.draw(hidden_name=self._main_image_name), position="free"
+            )
             # Top-level Layers panel (visible only when multiple layers exist).
             self.gui.register_ui_callback(self._ui_populate_layers_panel, position="panel")
 
@@ -576,6 +579,8 @@ def _filter_destroy(d: dict) -> dict:
 
         if getattr(self, "_image_logger", None) is not None:
             self._image_logger.clear_matching(owns)
+        if owns(getattr(self, "_main_image_name", "") or ""):
+            self._main_image_name = None
 
         # Drop example-registered side/free UI callbacks (panel/stats/rendering persist).
         if getattr(self, "gui", None) is not None:
@@ -1496,6 +1501,25 @@ def log_image(self, name: str, image: wp.array[Any] | np.ndarray) -> None:
         name = self._qualify(name)
         self._image_logger.log(name, image)
 
+    def log_main_image(self, name: str, image: wp.array[Any] | np.ndarray) -> None:
+        """Log an image and display it as the main viewer surface for this frame.
+
+        When a main image is logged for a frame, :class:`ViewerGL` skips the
+        normal 3D scene render and draws the image texture directly to the
+        window. If no main image is logged before a later :meth:`end_frame`,
+        the viewer returns to the normal 3D scene render for that frame.
+
+        Args:
+            name: Stable identifier for the image.
+            image: Image array. See :meth:`log_image` for accepted shapes and
+                dtypes.
+        """
+        if not isinstance(name, str) or not name:
+            raise ValueError("main image name must be a non-empty string")
+        name = self._qualify(name)
+        self._image_logger.log(name, image)
+        self._main_image_name = name
+
     @override
     def log_scalar(
         self,
@@ -1729,17 +1753,36 @@ def _update(self):
         if self.wind is not None:
             self.wind.update(dt)
 
-        # If the window was closed during event processing, skip rendering
-        if self.renderer.has_exit():
-            return
-
-        # Render the scene and present it
-        self.renderer.render(self.camera, self.objects, self.lines, self.wireframe_shapes, self.arrows)
+        try:
+            # If the window was closed during event processing, skip rendering
+            if self.renderer.has_exit():
+                return
+
+            # Render either the selected logged image or the 3D scene, then present it.
+            main_image_name = self._main_image_name
+            if main_image_name is not None:
+                texture = self._image_logger.get_texture(main_image_name)
+                if texture is None:
+                    self.renderer.render_texture(None, 0, 0)
+                else:
+                    self.renderer.render_texture(
+                        texture.texture_id,
+                        texture.texture_width,
+                        texture.texture_height,
+                        tile_count=texture.tile_count,
+                        tile_width=texture.tile_width,
+                        tile_height=texture.tile_height,
+                        atlas_cols=texture.atlas_cols,
+                    )
+            else:
+                self.renderer.render(self.camera, self.objects, self.lines, self.wireframe_shapes, self.arrows)
 
-        if self.gui:
-            self.gui.render_frame(update_fps=True)
+            if self.gui:
+                self.gui.render_frame(update_fps=True)
 
-        self.renderer.present()
+            self.renderer.present()
+        finally:
+            self._main_image_name = None
 
     def get_frame(self, target_image: wp.array | None = None, render_ui: bool = False) -> wp.array:
         """