diff --git a/python/mujoco/experimental/studio/window.cc b/python/mujoco/experimental/studio/window.cc new file mode 100644 index 00000000..e0b8507f --- /dev/null +++ b/python/mujoco/experimental/studio/window.cc @@ -0,0 +1,110 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static bool IsCuda() { +#ifdef CUDA + return true; +#else + return false; +#endif +} + +static bool IsCrd() { + const char* display = getenv("DISPLAY"); + return display ? strcmp(display, ":20") == 0 : false; +} + +static intptr_t GetImGuiContext() { + return reinterpret_cast(ImGui::GetCurrentContext()); +} + +static intptr_t GetImPlotContext() { + return reinterpret_cast(ImPlot::GetCurrentContext()); +} + +class Window { + public: + Window(const std::string& title, int width, int height, + std::string graphics_mode_str) { + pybind11::gil_scoped_release no_gil; + + mujoco::platform::RegisterResourceProviders(); + + mujoco::platform::Window::Config config; + config.gfx_mode = mujoco::platform::GraphicsModeFromString( + graphics_mode_str, mujoco::platform::GraphicsMode::FilamentOpenGl); + window_ = std::make_unique( + title, width, height, config); + ImPlot::CreateContext(); + } + + bool NewFrame() { + pybind11::gil_scoped_release no_gil; + const mujoco::platform::Window::Status window_status = window_->NewFrame(); + return window_status == mujoco::platform::Window::Status::kRunning; + } + + void Present(pybind11::bytes pixels) { + pybind11::gil_scoped_release no_gil; + std::string_view sv(pixels); + window_->EndFrame(); + window_->Present({(std::byte*)sv.data(), sv.size()}); + } + + int GetWidth() { return window_->GetWidth(); } + int GetHeight() { return window_->GetHeight(); } + + uint64_t GetNativeWindowHandle() { + return reinterpret_cast(window_->GetNativeWindowHandle()); + } + + std::string GetDropFile() { return window_->GetDropFile(); } + + private: + std::unique_ptr window_; +}; + +PYBIND11_MODULE(window, m, pybind11::mod_gil_not_used()) { + m.def("IsCrd", &IsCrd); + m.def("IsCuda", &IsCuda); + m.def("GetImGuiContext", &GetImGuiContext); + m.def("GetImPlotContext", &GetImPlotContext); + + pybind11::class_ cls(m, "Window"); + cls.def(pybind11::init()); + cls.def("NewFrame", &Window::NewFrame); + cls.def("Present", &Window::Present); + cls.def("GetWidth", &Window::GetWidth); + cls.def("GetHeight", &Window::GetHeight); + cls.def("GetDropFile", &Window::GetDropFile); + cls.def("GetNativeWindowHandle", &Window::GetNativeWindowHandle); +}