diff --git a/python/mujoco/simulate.cc b/python/mujoco/simulate.cc index 56fe4437..336ffb71 100644 --- a/python/mujoco/simulate.cc +++ b/python/mujoco/simulate.cc @@ -13,17 +13,22 @@ // limitations under the License. #include -#include +#include // NOLINT(build/c++11) #include #include +#include #include -#include +#include // NOLINT(build/c++11) +#include #include +#include #include #include #include +#include #include "errors.h" +#include "indexers.h" #include "structs.h" #include #include @@ -91,6 +96,7 @@ class SimulateWrapper { void Destroy() { if (simulate_) { + ClearImages(); delete simulate_; simulate_ = nullptr; destroyed_.store(1); @@ -148,6 +154,60 @@ class SimulateWrapper { void ClearFigures() { simulate_->user_figures_.clear(); } + void SetTexts( + const std::vector>& + texts) { + // Collection of [font, gridpos, text1, text2] tuples for overlay text + std::vector> user_texts; + for (const auto& [font, gridpos, text1, text2] : texts) { + user_texts.push_back(std::make_tuple(font, gridpos, text1, text2)); + } + + // Set them all at once to prevent text flickering. + simulate_->user_texts_ = user_texts; + } + + void ClearTexts() { simulate_->user_texts_.clear(); } + + void SetImages( + const std::vector> viewports_images + ) { + // Clear previous images to prevent memory leaks + ClearImages(); + + for (const auto& [viewport, image] : viewports_images) { + auto buf = image.request(); + if (buf.ndim != 3) { + throw std::invalid_argument("image must have 3 dimensions (H, W, C)"); + } + if (static_cast(buf.shape[2]) != 3) { + throw std::invalid_argument("image must have 3 channels"); + } + if (buf.itemsize != sizeof(unsigned char)) { + throw std::invalid_argument("image must be uint8 format"); + } + + // 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() { + // Free memory for each image before clearing the vector + for (const auto& [viewport, image_ptr] : simulate_->user_images_) { + delete[] image_ptr; + } + simulate_->user_images_.clear(); + } + private: mujoco::Simulate* simulate_; std::atomic_int destroyed_ = 0; @@ -249,6 +309,11 @@ PYBIND11_MODULE(_simulate, pymodule) { .def("set_figures", &SimulateWrapper::SetFigures, py::arg("viewports_figures")) .def("clear_figures", &SimulateWrapper::ClearFigures) + .def("set_texts", &SimulateWrapper::SetTexts, py::arg("overlay_texts")) + .def("clear_texts", &SimulateWrapper::ClearTexts) + .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) diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py index 65852c87..c1de94ec 100644 --- a/python/mujoco/viewer.py +++ b/python/mujoco/viewer.py @@ -23,7 +23,7 @@ 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 @@ -115,9 +115,22 @@ class Handle: return sim.viewport return None - def set_figures(self, viewports_figures): + def set_figures( + self, viewports_figures: Union[Tuple[mujoco.MjrRect, mujoco.MjvFigure], + List[Tuple[mujoco.MjrRect, mujoco.MjvFigure]]] + ): + """Overlay figures on the viewer. + + Args: + viewports_figures: Single tuple or list of tuples of (viewport, figure) + viewport: Rectangle defining position and size of the figure + figure: MjvFigure object containing the figure data to display + """ sim = self._sim() if sim is not None: + # Convert single tuple to list if needed + if isinstance(viewports_figures, tuple): + viewports_figures = [viewports_figures] sim.set_figures(viewports_figures) def clear_figures(self): @@ -125,6 +138,73 @@ class Handle: if sim is not None: sim.clear_figures() + def set_texts(self, texts: Union[Tuple[Optional[int], Optional[int], Optional[str], Optional[str]], + List[Tuple[Optional[int], Optional[int], Optional[str], Optional[str]]]]): + """Overlay text on the viewer. + + Args: + texts: Single tuple or list of tuples of (font, gridpos, text1, text2) + font: Font style from mujoco.mjtFontScale + gridpos: Position of text box from mujoco.mjtGridPos + text1: Left text column, defaults to empty string if None + text2: Right text column, defaults to empty string if None + """ + sim = self._sim() + if sim is not None: + # Convert single tuple to list if needed + if isinstance(texts, tuple): + texts = [texts] + + # Convert None values to empty strings + default_font = mujoco.mjtFontScale.mjFONTSCALE_150 + default_gridpos = mujoco.mjtGridPos.mjGRID_TOPLEFT + processed_texts = [( + default_font if font is None else font, + default_gridpos if gridpos is None else gridpos, + "" if text1 is None else text1, + "" if text2 is None else text2) + for font, gridpos, text1, text2 in texts] + + sim.set_texts(processed_texts) + + def clear_texts(self): + sim = self._sim() + if sim is not None: + sim.clear_texts() + + def set_images( + self, viewports_images: Union[Tuple[mujoco.MjrRect, np.ndarray], + List[Tuple[mujoco.MjrRect, np.ndarray]]] + ): + """Overlay images on the viewer. + + Args: + viewports_images: Single tuple or list of tuples of (viewport, image) + viewport: Rectangle defining position and size of the image + image: RGB image with shape (height, width, 3) + """ + sim = self._sim() + if sim is not None: + # Convert single tuple to list if needed + if isinstance(viewports_images, tuple): + viewports_images = [viewports_images] + + processed_images = [] + for viewport, image in viewports_images: + targ_shape = (viewport.height, viewport.width) + # Check if image is already the correct shape + if image.shape[:2] != targ_shape: + raise ValueError(f"Image shape {image.shape[:2]} does not match target shape {targ_shape}") + flipped = np.flip(image, axis=0) + contiguous = np.ascontiguousarray(flipped) + processed_images.append((viewport, contiguous)) + sim.set_images(processed_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: diff --git a/simulate/simulate.cc b/simulate/simulate.cc index 2f5ee146..8d5013a3 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -464,7 +464,7 @@ void InitializeSensor(mj::Simulate* sim) { // title mju::strcpy_arr(figsensor.title, "Sensor data"); - // y-tick nubmer format + // y-tick number format mju::strcpy_arr(figsensor.yformat, "%.1f"); // grid size @@ -546,6 +546,16 @@ void ShowFigure(mj::Simulate* sim, mjrRect viewport, mjvFigure* fig){ mjr_figure(viewport, fig, &sim->platform_ui->mjr_context()); } +void ShowOverlayText(mj::Simulate* sim, mjrRect viewport, int font, int gridpos, + std::string text1, std::string text2) { + mjr_overlay(font, gridpos, viewport, text1.c_str(), text2.c_str(), + &sim->platform_ui->mjr_context()); +} + +void ShowImage(mj::Simulate* sim, mjrRect viewport, const unsigned char* image) { + mjr_drawPixels(image, nullptr, viewport, &sim->platform_ui->mjr_context()); +} + // load state from history buffer static void LoadScrubState(mj::Simulate* sim) { // get index into circular buffer @@ -2656,6 +2666,16 @@ void Simulate::Render() { ShowFigure(this, viewport, &figure); } + // overlay text + for (auto& [font, gridpos, text1, text2] : this->user_texts_) { + ShowOverlayText(this, rect, font, gridpos, text1, text2); + } + + // user images + for (auto& [viewport, image] : this->user_images_) { + ShowImage(this, viewport, image); + } + // finalize this->platform_ui->SwapBuffers(); } diff --git a/simulate/simulate.h b/simulate/simulate.h index 23bdb604..717d1871 100644 --- a/simulate/simulate.h +++ b/simulate/simulate.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include @@ -250,10 +251,12 @@ class Simulate { mjvFigure figsize = {}; mjvFigure figsensor = {}; - // additional user-defined visualization geoms (used in passive mode) + // additional user-defined visualization mjvScene* user_scn = nullptr; mjtByte user_scn_flags_prev_[mjNRNDFLAG]; std::vector> user_figures_; + std::vector> user_texts_; + std::vector> user_images_; // OpenGL rendering and UI int refresh_rate = 60;