add handles for text and image overlays

python/mujoco/simulate.cc

@@ -148,6 +148,50 @@ class SimulateWrapper {
 
   void ClearFigures() { simulate_->user_figures_.clear(); }
 
+  void SetOverlayText(
+      const std::vector<std::tuple<int, int, std::string, std::string>>& overlay_texts) {
+    // Collection of [font, gridpos, text1, text2] tuples for overlay text
+    std::vector<std::tuple<int, int, std::string, std::string>> user_overlay_text;
+    for (const auto& [font, gridpos, text1, text2] : overlay_texts) {
+      user_overlay_text.push_back(std::make_tuple(font, gridpos, text1, text2));
+    }
+
+    // Set them all at once to prevent overlay text flickering.
+    simulate_->user_text_ = user_overlay_text;
+  }
+
+  void ClearOverlayText() { simulate_->user_text_.clear(); }
+
+  void SetImages(
+    const std::vector<std::tuple<mjrRect, pybind11::array_t<unsigned char>>>& viewport_images
+  ) {
+    // Clear previous images to prevent memory leaks
+    simulate_->user_images_.clear();
+    
+    for (const auto& [viewport, image] : viewport_images) {
+      auto buf = image.request();
+      if (static_cast<int>(buf.shape[2]) != 3) {
+        throw std::invalid_argument("image must have 3 channels");
+      }
+      if (buf.ndim != 3) {
+        throw std::invalid_argument("image must have 3 dimensions (H, W, C)");
+      }
+      
+      // Calculate size of the image data
+      size_t height = buf.shape[0];
+      size_t width = buf.shape[1];
+      size_t size = height * width * 3;
+      
+      // Make a copy of the image data to prevent flickering
+      unsigned char* image_copy = new unsigned char[size];
+      std::memcpy(image_copy, buf.ptr, size);
+      
+      simulate_->user_images_.push_back(std::make_tuple(viewport, image_copy));
+    }
+  }
+
+  void ClearImages() { simulate_->user_images_.clear(); }
+
  private:
   mujoco::Simulate* simulate_;
   std::atomic_int destroyed_ = 0;
@@ -249,6 +293,12 @@ PYBIND11_MODULE(_simulate, pymodule) {
       .def("set_figures", &SimulateWrapper::SetFigures,
            py::arg("viewports_figures"))
       .def("clear_figures", &SimulateWrapper::ClearFigures)
+      .def("overlay_text", &SimulateWrapper::SetOverlayText,
+           py::arg("overlay_texts"))
+      .def("clear_overlay_text", &SimulateWrapper::ClearOverlayText)
+      .def("set_images", &SimulateWrapper::SetImages,
+           py::arg("viewports_images"))
+      .def("clear_images", &SimulateWrapper::ClearImages)
       .def_property_readonly("m", &SimulateWrapper::GetModel)
       .def_property_readonly("d", &SimulateWrapper::GetData)
       .def_property_readonly("viewport", &SimulateWrapper::GetViewport)

python/mujoco/viewer.py

@@ -23,9 +23,8 @@ import queue
 import sys
 import threading
 import time
-from typing import Callable, Optional, Tuple, Union
+from typing import Callable, List, Optional, Tuple, Union
 import weakref
-
 import glfw
 import mujoco
 from mujoco import _simulate
@@ -115,7 +114,7 @@ class Handle:
       return sim.viewport
     return None
 
-  def set_figures(self, viewports_figures):
+  def set_figures(self, viewports_figures: List[Tuple[mujoco.MjrRect, mujoco.MjvFigure]]):
     sim = self._sim()
     if sim is not None:
       sim.set_figures(viewports_figures)
@@ -125,6 +124,45 @@ class Handle:
     if sim is not None:
       sim.clear_figures()
 
+  def overlay_text(self, overlay_texts: List[Tuple[int, int, str, str]]):
+    """ Overlay text on the viewer.
+    
+    Args:
+      overlay_texts: List of tuples of (font, gridpos, text1, text2)
+      let:
+        font: Font style from mujoco.mjtFontScale
+        gridpos: Position of text box from mujoco.mjtGridPos
+        text1: Left text column
+        text2: Right text column
+    """
+    sim = self._sim()
+    if sim is not None:
+      sim.overlay_text(overlay_texts)
+  
+  def clear_overlay_text(self):
+    sim = self._sim()
+    if sim is not None:
+      sim.clear_overlay_text()
+
+  def set_images(self, viewports_images: List[Tuple[mujoco.MjrRect, np.ndarray]]):
+    sim = self._sim()
+    if sim is not None:
+      # Nearest neighbor resize
+      resize = lambda a, s: a[(np.arange(s[0]) * a.shape[0]) // s[0]][:, (np.arange(s[1]) * a.shape[1]) // s[1]]
+      resized_viewports_images = []
+      for viewport, image in viewports_images:
+        targ_shape = (viewport.height, viewport.width)
+        resized = resize(image, targ_shape)
+        resized = np.flip(resized, axis=0)
+        resized = np.ascontiguousarray(resized)
+        resized_viewports_images.append((viewport, resized))
+      sim.set_images(resized_viewports_images)
+
+  def clear_images(self):
+    sim = self._sim()
+    if sim is not None:
+      sim.clear_images()
+
   def close(self):
     sim = self._sim()
     if sim is not None: