Correctly handle segmentation and ID_COLOR flag in filament.

PiperOrigin-RevId: 846319715
Change-Id: I244ae52c34079b79d3179fafdf6aa506cc89e1ff

src/experimental/filament/filament/drawable.cc

@@ -161,7 +161,7 @@ void Drawable::Update(const mjModel* model, const mjvScene* scene,
   }
 
   SetTransform(geom);
-  UpdateMaterial(geom);
+  UpdateMaterial(geom, scene->flags[mjRND_IDCOLOR]);
 }
 
 void Drawable::AddMesh(int data_id) {
@@ -203,11 +203,6 @@ 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) {
@@ -321,7 +316,7 @@ void Drawable::SetTransform(const mjvGeom& geom) {
   }
 }
 
-void Drawable::UpdateMaterial(const mjvGeom& geom) {
+void Drawable::UpdateMaterial(const mjvGeom& geom, bool use_segid_color) {
   ObjectManager* object_mgr = material_.GetObjectManager();
   const mjModel* model = object_mgr->GetModel();
 
@@ -405,7 +400,6 @@ 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];
@@ -414,8 +408,8 @@ void Drawable::UpdateMaterial(const mjvGeom& geom) {
   }
 
   if (geom.segid >= 0) {
-    uint32_t segmentation_color = geom.segid;
+    uint32_t segmentation_color = geom.segid + 1;
-    if (use_distinct_segmentation_colors_) {
+    if (!use_segid_color) {
       constexpr double phi1 = 1.61803398874989484820;  // Cached Phi(1).
       constexpr double coef1 = 1.0 / phi1;
       const double index = static_cast<double>(geom.segid);

src/experimental/filament/filament/drawable.h

@@ -53,9 +53,6 @@ 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);
@@ -66,11 +63,10 @@ class Drawable {
   void SetTransform(const mjvGeom& geom);
 
   // Updates the material parameters of the drawable for rendering.
-  void UpdateMaterial(const mjvGeom& geom);
+  void UpdateMaterial(const mjvGeom& geom, bool use_segid_color);
 
   Material material_;
   Renderables renderables_;
-  bool use_distinct_segmentation_colors_ = false;
 };
 
 }  // namespace mujoco

src/experimental/filament/filament/filament_context.cc

@@ -115,10 +115,7 @@ FilamentContext::FilamentContext(const mjrFilamentConfig* config,
     }
   }
 
-  scene_view_ = std::make_unique<SceneView>(
+  scene_view_ = std::make_unique<SceneView>(engine_, object_manager_.get());
-      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,7 +60,6 @@ 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,11 +210,6 @@ 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);
@@ -359,8 +354,6 @@ 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,10 +55,6 @@ 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);
 
@@ -104,7 +100,6 @@ 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,10 +51,6 @@ 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

@@ -142,7 +142,6 @@ 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)