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

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),
+    : material_(object_mgr), renderables_(object_mgr->GetEngine()) {
-      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.roughness =
-    textures.occlusion = object_mgr->GetTexture(geom.matid, mjTEXROLE_OCCLUSION);
+        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)
+    // `data_id` is the id of the mesh in model (i.e. the geom has mesh
-    // and `mesh_texcoordadr` stores the address of the mesh uvs if it has them.
+    // 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).
+    uint32_t segmentation_color = geom.segid;
-    constexpr double coef1 = 1.0 / phi1;
+    if (use_distinct_segmentation_colors_) {
-    const double index = static_cast<double>(geom.segid);
+      constexpr double phi1 = 1.61803398874989484820;  // Cached Phi(1).
-    const double sample = std::fmod(0.5 + coef1 * index, 1.0);
+      constexpr double coef1 = 1.0 / phi1;
-    uint32_t segmentation_color = 0x01000000 * sample;
+      const double index = static_cast<double>(geom.segid);
-    const uint8_t red = (segmentation_color >> 16) & 0xff;
+      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<float>(red) / 255.0f;
     params.segmentation_color.y = static_cast<float>(green) / 255.0f;
     params.segmentation_color.z = static_cast<float>(blue) / 255.0f;

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

src/experimental/filament/filament/filament_context.cc

@@ -115,7 +115,10 @@ FilamentContext::FilamentContext(const mjrFilamentConfig* config,
     }
   }
 
-  scene_view_ = std::make_unique<SceneView>(engine_, object_manager_.get());
+  scene_view_ = std::make_unique<SceneView>(
+      engine_, object_manager_.get());
+  scene_view_->SetUseDistinctSegmentationColors(
+      config_.use_distinct_segmentation_colors);
   if (config_.enable_gui) {
     gui_view_ = std::make_unique<GuiView>(engine_, object_manager_.get());
   }

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);

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<Drawable>(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));
   }

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

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);

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)