From a942fd48759705eb81765442c05b2f9cdbc8f048 Mon Sep 17 00:00:00 2001 From: Sam Haves Date: Wed, 3 Dec 2025 19:25:29 -0800 Subject: [PATCH] Make filament segmentation shader distinct colors optional. Previously we were computing at material update time a color for the geom based of it's segid. This made segids unrecoverable from the resulting buffer. With this change we allow SceneView creators to specify if they want the segmentation buffer to be purely for visualization with distinct colors. PiperOrigin-RevId: 840016870 Change-Id: I5920f248ace0901023e6545c7d8dd2b23f9aa397 --- .../filament/filament/drawable.cc | 42 ++++++++++++------- src/experimental/filament/filament/drawable.h | 5 +++ .../filament/filament/filament_context.cc | 5 ++- src/experimental/filament/filament/material.h | 1 + .../filament/filament/scene_view.cc | 7 ++++ .../filament/filament/scene_view.h | 5 +++ .../filament/render_context_filament.h | 4 ++ src/experimental/studio/app.cc | 1 + 8 files changed, 54 insertions(+), 16 deletions(-) diff --git a/src/experimental/filament/filament/drawable.cc b/src/experimental/filament/filament/drawable.cc index a44b4d9b..3fc52200 100644 --- a/src/experimental/filament/filament/drawable.cc +++ b/src/experimental/filament/filament/drawable.cc @@ -70,13 +70,12 @@ static constexpr int kArrow2TopConeDisk = 3; static constexpr int kArrow2BottomConeDisk = 4; Drawable::Drawable(ObjectManager* object_mgr, const mjvGeom& geom) - : material_(object_mgr), - renderables_(object_mgr->GetEngine()) { + : material_(object_mgr), renderables_(object_mgr->GetEngine()) { if (geom.category == mjCAT_DECOR) { renderables_.DisableShadows(); } - switch ((mjtGeom)geom.type) { + switch ((mjtGeom)geom.type) { case mjGEOM_MESH: AddMesh(geom.dataid); break; @@ -204,6 +203,11 @@ void Drawable::SetDrawMode(Material::DrawMode mode) { renderables_.SetMaterialInstance(material_.GetMaterialInstance(mode)); } +void Drawable::SetUseDistinctSegmentationColors( + bool use_distinct_segmentation_colors) { + use_distinct_segmentation_colors_ = use_distinct_segmentation_colors; +} + void Drawable::SetTransform(const mjvGeom& geom) { // Flex and skin geometries are in global space. if (geom.type == mjGEOM_FLEX || geom.type == mjGEOM_SKIN) { @@ -328,8 +332,10 @@ void Drawable::UpdateMaterial(const mjvGeom& geom) { textures.emissive = object_mgr->GetTexture(geom.matid, mjTEXROLE_EMISSIVE); textures.orm = object_mgr->GetTexture(geom.matid, mjTEXROLE_ORM); textures.metallic = object_mgr->GetTexture(geom.matid, mjTEXROLE_METALLIC); - textures.roughness = object_mgr->GetTexture(geom.matid, mjTEXROLE_ROUGHNESS); - textures.occlusion = object_mgr->GetTexture(geom.matid, mjTEXROLE_OCCLUSION); + textures.roughness = + object_mgr->GetTexture(geom.matid, mjTEXROLE_ROUGHNESS); + textures.occlusion = + object_mgr->GetTexture(geom.matid, mjTEXROLE_OCCLUSION); material_.UpdateTextures(textures); } @@ -351,8 +357,9 @@ void Drawable::UpdateMaterial(const mjvGeom& geom) { } // Check to see if we're dealing with a mesh with texture coordinates. - // `data_id` is the id of the mesh in model (i.e. the geom has mesh geometry) - // and `mesh_texcoordadr` stores the address of the mesh uvs if it has them. + // `data_id` is the id of the mesh in model (i.e. the geom has mesh + // geometry) and `mesh_texcoordadr` stores the address of the mesh uvs if it + // has them. bool has_texcoords = false; if ((geom.type == mjGEOM_MESH || geom.type == mjGEOM_SDF) && geom.dataid >= 0 && model->mesh_texcoordadr[geom.dataid / 2] >= 0) { @@ -366,7 +373,7 @@ void Drawable::UpdateMaterial(const mjvGeom& geom) { material_.SetNormalMaterialType(ObjectManager::kPhongColor); } } else if (textures.color->getTarget() == - filament::Texture::Sampler::SAMPLER_CUBEMAP) { + filament::Texture::Sampler::SAMPLER_CUBEMAP) { if (geom.rgba[3] < 1.0f) { material_.SetNormalMaterialType(ObjectManager::kPhongCubeFade); } else { @@ -392,6 +399,7 @@ void Drawable::UpdateMaterial(const mjvGeom& geom) { params.emissive = geom.emission; params.specular = geom.specular; params.glossiness = geom.shininess; + params.use_distinct_segmentation_colors = use_distinct_segmentation_colors_; if (geom.matid >= 0) { params.metallic = model->mat_metallic[geom.matid]; params.roughness = model->mat_roughness[geom.matid]; @@ -400,14 +408,18 @@ void Drawable::UpdateMaterial(const mjvGeom& geom) { } if (geom.segid >= 0) { - constexpr double phi1 = 1.61803398874989484820; // Cached Phi(1). - constexpr double coef1 = 1.0 / phi1; - const double index = static_cast(geom.segid); - const double sample = std::fmod(0.5 + coef1 * index, 1.0); - uint32_t segmentation_color = 0x01000000 * sample; - const uint8_t red = (segmentation_color >> 16) & 0xff; + uint32_t segmentation_color = geom.segid; + if (use_distinct_segmentation_colors_) { + constexpr double phi1 = 1.61803398874989484820; // Cached Phi(1). + constexpr double coef1 = 1.0 / phi1; + const double index = static_cast(geom.segid); + const double sample = std::fmod(0.5 + coef1 * index, 1.0); + segmentation_color = 0x01000000 * sample; + } + + const uint8_t red = (segmentation_color >> 0) & 0xff; const uint8_t green = (segmentation_color >> 8) & 0xff; - const uint8_t blue = (segmentation_color >> 0) & 0xff; + const uint8_t blue = (segmentation_color >> 16) & 0xff; params.segmentation_color.x = static_cast(red) / 255.0f; params.segmentation_color.y = static_cast(green) / 255.0f; params.segmentation_color.z = static_cast(blue) / 255.0f; diff --git a/src/experimental/filament/filament/drawable.h b/src/experimental/filament/filament/drawable.h index c44c37e4..2b2ffb3c 100644 --- a/src/experimental/filament/filament/drawable.h +++ b/src/experimental/filament/filament/drawable.h @@ -53,6 +53,10 @@ class Drawable { // beginFrame/endFrame. void SetDrawMode(Material::DrawMode mode); + // Updates whether to remap segmentation IDs to distinct colors when rendering + // segmentation. This is only really useful for visualization purposes. + void SetUseDistinctSegmentationColors(bool use_distinct_segmentation_colors); + private: void AddMesh(int data_id); void AddHeightField(int hfield_id); @@ -66,6 +70,7 @@ class Drawable { Material material_; Renderables renderables_; + bool use_distinct_segmentation_colors_ = false; }; } // namespace mujoco diff --git a/src/experimental/filament/filament/filament_context.cc b/src/experimental/filament/filament/filament_context.cc index e06e8b42..7ffd22e3 100644 --- a/src/experimental/filament/filament/filament_context.cc +++ b/src/experimental/filament/filament/filament_context.cc @@ -115,7 +115,10 @@ FilamentContext::FilamentContext(const mjrFilamentConfig* config, } } - scene_view_ = std::make_unique(engine_, object_manager_.get()); + scene_view_ = std::make_unique( + engine_, object_manager_.get()); + scene_view_->SetUseDistinctSegmentationColors( + config_.use_distinct_segmentation_colors); if (config_.enable_gui) { gui_view_ = std::make_unique(engine_, object_manager_.get()); } diff --git a/src/experimental/filament/filament/material.h b/src/experimental/filament/filament/material.h index c21f9fb1..6313b5e5 100644 --- a/src/experimental/filament/filament/material.h +++ b/src/experimental/filament/filament/material.h @@ -60,6 +60,7 @@ class Material { float roughness = -1.0f; float emissive = -1.0f; bool tex_uniform = false; + bool use_distinct_segmentation_colors = false; }; Material(ObjectManager* object_mgr); diff --git a/src/experimental/filament/filament/scene_view.cc b/src/experimental/filament/filament/scene_view.cc index 4977bf47..f738f284 100644 --- a/src/experimental/filament/filament/scene_view.cc +++ b/src/experimental/filament/filament/scene_view.cc @@ -210,6 +210,11 @@ void SceneView::SetColorGradingOptions(const ColorGradingOptions& opts) { color_grading_options_ = opts; } +void SceneView::SetUseDistinctSegmentationColors( + bool use_distinct_segmentation_colors) { + use_distinct_segmentation_colors_ = use_distinct_segmentation_colors; +} + void SceneView::SetEnvironmentLight(std::string_view filename, float intensity) { auto* ibl = object_mgr_->LoadFallbackIndirectLight(filename, intensity); @@ -349,6 +354,8 @@ void SceneView::UpdateScene(const mjrContext* context, const mjvScene* scene) { auto drawable = std::make_unique(object_mgr_, *geom); drawable->AddToScene(scene_); + drawable->SetUseDistinctSegmentationColors( + use_distinct_segmentation_colors_); drawable->Update(object_mgr_->GetModel(), scene, *geom); drawables_.push_back(std::move(drawable)); } diff --git a/src/experimental/filament/filament/scene_view.h b/src/experimental/filament/filament/scene_view.h index 223e4940..596061aa 100644 --- a/src/experimental/filament/filament/scene_view.h +++ b/src/experimental/filament/filament/scene_view.h @@ -55,6 +55,10 @@ class SceneView { // Updates the color grading options for the main render view. void SetColorGradingOptions(const ColorGradingOptions& opts); + // Updates whether to remap segmentation IDs to distinct colors when rendering + // segmentation. This is only really useful for visualization purposes. + void SetUseDistinctSegmentationColors(bool use_distinct_segmentation_colors); + // Updates the environment light using the KTX image at the given path. void SetEnvironmentLight(std::string_view filename, float intensity); @@ -100,6 +104,7 @@ class SceneView { ColorGradingOptions color_grading_options_; DrawMode active_mode_ = DrawMode::kNumDrawModes; float aspect_ratio_ = 1.0f; + bool use_distinct_segmentation_colors_ = false; }; } // namespace mujoco diff --git a/src/experimental/filament/render_context_filament.h b/src/experimental/filament/render_context_filament.h index 6d33b655..0d215f9f 100644 --- a/src/experimental/filament/render_context_filament.h +++ b/src/experimental/filament/render_context_filament.h @@ -51,6 +51,10 @@ struct mjrFilamentConfig { // Whether or not to enable GUI rendering. bool enable_gui; + + // Whether to remap segmentation IDs to distinct colors for visual display. + // This is only really useful for visualization purposes in studio. + bool use_distinct_segmentation_colors; }; void mjr_defaultFilamentConfig(mjrFilamentConfig* config); diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc index a78fde75..fc44e2f1 100644 --- a/src/experimental/studio/app.cc +++ b/src/experimental/studio/app.cc @@ -164,6 +164,7 @@ App::App(int width, int height, std::string ini_path, render_config.load_asset = &App::LoadAssetCallback; render_config.load_asset_user_data = this; render_config.enable_gui = true; + render_config.use_distinct_segmentation_colors = true; #if defined(USE_FILAMENT_OPENGL) render_config.graphics_api = mjGFX_OPENGL; #elif defined(USE_FILAMENT_VULKAN)