Move segmentation coloring logic into renderer.

PiperOrigin-RevId: 927300828
Change-Id: I994c5742ec30993d2245df167042e72820f91ac7

src/experimental/filament/compat/scene_geom_util.cc

@@ -183,20 +183,7 @@ static void UpdateGeomMaterial(mjrRenderable* renderable, const mjvGeom& geom,
     material.roughness = model->mat_roughness[geom.matid];
   }
 
-  if (geom.segid >= 0) {
+  material.segmentation_id = geom.segid;
-    uint32_t segmentation_color = geom.segid + 1;
-    const bool use_segid_color = render_flags[mjRND_IDCOLOR];
-    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);
-      const double sample = std::fmod(0.5 + coef1 * index, 1.0);
-      segmentation_color = 0x01000000 * sample;
-    }
-    material.segmentation_color[0] = (segmentation_color >> 0);
-    material.segmentation_color[1] = (segmentation_color >> 8);
-    material.segmentation_color[2] = (segmentation_color >> 16);
-  }
 
   // UvScale only applies to objects that don't have explicit UV coordinates
   // in their vertex buffer. Instead, we set the UV coordinate to be the same

src/experimental/filament/filament/material_manager.cc

@@ -195,7 +195,9 @@ MaterialManager::MaterialKey MaterialManager::PrepareMaterialInstance(
   ObjectManager::MaterialType type;
   if (draw_mode == mjDRAW_MODE_DEPTH) {
     type = ObjectManager::kUnlitDepth;
-  } else if (draw_mode == mjDRAW_MODE_SEGMENTATION) {
+  } else if (draw_mode == mjDRAW_MODE_SEGMENTATION_BY_ID) {
+    type = ObjectManager::kUnlitSegmentation;
+  } else if (draw_mode == mjDRAW_MODE_SEGMENTATION_BY_COLOR) {
     type = ObjectManager::kUnlitSegmentation;
   } else {
     type = GetMaterialType(material, mesh);
@@ -238,10 +240,12 @@ void MaterialManager::UpdateMaterialInstance(filament::MaterialInstance* instanc
                            ReadFloat4(material.color));
   }
   if (fmaterial->hasParameter("SegmentationColor")) {
-    filament::math::float4 color{
+    const uint8_t r = material.segmentation_id >> 0;
-        static_cast<float>(material.segmentation_color[0]) / 255.0f,
+    const uint8_t g = material.segmentation_id >> 8;
-        static_cast<float>(material.segmentation_color[1]) / 255.0f,
+    const uint8_t b = material.segmentation_id >> 16;
-        static_cast<float>(material.segmentation_color[2]) / 255.0f, 1.0f};
+    const filament::math::float4 color{static_cast<float>(r) / 255.0f,
+                                       static_cast<float>(g) / 255.0f,
+                                       static_cast<float>(b) / 255.0f, 1.0f};
     instance->setParameter("SegmentationColor", filament::RgbaType::LINEAR,
                            color);
   }

src/experimental/filament/filament/renderable.cc

@@ -14,6 +14,7 @@
 
 #include "experimental/filament/filament/renderable.h"
 
+#include <cmath>
 #include <cstdint>
 #include <functional>
 #include <numbers>
@@ -257,15 +258,19 @@ void Renderable::Prepare(std::span<const mjrRenderRequest*> requests,
       }
     }
 
-    if (request->draw_mode != mjDRAW_MODE_SEGMENTATION) {
+    if (request->draw_mode == mjDRAW_MODE_SEGMENTATION_BY_COLOR) {
+      constexpr double phi1 = 1.61803398874989484820;  // Cached Phi(1).
+      constexpr double coef1 = 1.0 / phi1;
+      const double index = static_cast<double>(material.segmentation_id);
+      const double sample = std::fmod(0.5 + coef1 * index, 1.0);
+      material.segmentation_id = 0x01000000 * sample;
+    } else if (request->draw_mode != mjDRAW_MODE_SEGMENTATION_BY_ID) {
       // Clear out the segmentation color in order to take advantage of shared
       // materials.
-      material.segmentation_color[0] = 0;
+      material.segmentation_id = 0;
-      material.segmentation_color[1] = 0;
-      material.segmentation_color[2] = 0;
     }
 
-    const bool reflective = request->draw_mode == mjDRAW_MODE_COLOR &&
+    const bool reflective = request->draw_mode == mjDRAW_MODE_DEFAULT &&
                             request->enable_reflections &&
                             material.reflectance > 0.0;
     if (reflective) {

src/experimental/filament/filament/scene_view.cc

@@ -240,7 +240,8 @@ void SceneView::Render(filament::Renderer* renderer, const mjrRenderRequest& req
                               request.viewport.width, request.viewport.height);
   filament::View* view = main_view_;
   if (request.draw_mode == mjDRAW_MODE_DEPTH ||
-      request.draw_mode == mjDRAW_MODE_SEGMENTATION) {
+      request.draw_mode == mjDRAW_MODE_SEGMENTATION_BY_ID ||
+      request.draw_mode == mjDRAW_MODE_SEGMENTATION_BY_COLOR) {
     view = depth_segment_view_;
   }
   view->setViewport(viewport);

src/experimental/filament/mjr_compat.cc

@@ -58,7 +58,7 @@ class CompatContext {
   }
 
  private:
-  mjrDrawMode draw_mode_ = mjDRAW_MODE_COLOR;
+  mjrDrawMode draw_mode_ = mjDRAW_MODE_DEFAULT;
   UniquePtr<mjrfContext> context_;
   std::unique_ptr<SceneBridge> scene_bridge_;
   mjtFramebuffer framebuffer_ = mjFB_WINDOW;
@@ -74,13 +74,17 @@ void CompatContext::Render(const mjrRect& viewport, const mjvScene* scene) {
   scene_bridge_->Update(viewport, scene);
 
   if (scene->flags[mjRND_SEGMENT]) {
-    draw_mode_ = mjDRAW_MODE_SEGMENTATION;
+    if (scene->flags[mjRND_IDCOLOR]) {
+      draw_mode_ = mjDRAW_MODE_SEGMENTATION_BY_ID;
+    } else {
+      draw_mode_ = mjDRAW_MODE_SEGMENTATION_BY_COLOR;
+    }
   } else if (scene->flags[mjRND_DEPTH]) {
     draw_mode_ = mjDRAW_MODE_DEPTH;
   } else if (scene->flags[mjRND_WIREFRAME]) {
     draw_mode_ = mjDRAW_MODE_WIREFRAME;
   } else {
-    draw_mode_ = mjDRAW_MODE_COLOR;
+    draw_mode_ = mjDRAW_MODE_DEFAULT;
   }
 
   if (framebuffer_ == mjFB_WINDOW) {

src/experimental/filament/render_context_filament.cc

@@ -81,9 +81,6 @@ void mjr_defaultMaterial(mjrMaterial* material) {
   memset(material, 0, sizeof(mjrMaterial));
   setf(material->color, {1.f, 1.f, 1.f, 1.f});
   setf(material->uv_scale, {1, 1, 1});
-  material->segmentation_color[0] = 255;
-  material->segmentation_color[1] = 255;
-  material->segmentation_color[2] = 255;
   material->emissive = -1.0f;
   material->specular = -1.0f;
   material->glossiness = -1.0f;

src/experimental/filament/render_context_filament.h

@@ -127,15 +127,16 @@ void mjrf_destroyContext(mjrfContext* ctx);
 
 // Describes the look/intention of the final rendered image.
 typedef enum mjrDrawMode_ {
-  // Render the scene with "normal" colors and lighting.
+  // Render the scene with default settings for colors and lighting (shading).
-  mjDRAW_MODE_COLOR,
+  mjDRAW_MODE_DEFAULT,
   // Render the scene as a wireframe.
   mjDRAW_MODE_WIREFRAME,
   // Render the scene as a grayscale depth map.
   mjDRAW_MODE_DEPTH,
-  // Render each object with a unique, uniform (flat) color regardless of
+  // Render each object using its segmentation id (see mjrMaterial).
-  // lighting and texture.
+  mjDRAW_MODE_SEGMENTATION_BY_ID,
-  mjDRAW_MODE_SEGMENTATION,
+  // Render each object with a unique color based on its segmentation id.
+  mjDRAW_MODE_SEGMENTATION_BY_COLOR,
 } mjrDrawMode;
 
 // Parameters describing the camera to use for rendering an image.
@@ -595,8 +596,9 @@ struct mjrMaterial {
   // The color of the object. Defaults to white.
   float color[4];
 
-  // The color to use for segmentation rendering. Defaults to white.
+  // The ID to use for segmentation rendering. These IDs will be mapped to a
-  mjtByte segmentation_color[3];
+  // RGB8 color, so only the first 24 bits are used.
+  int32_t segmentation_id;
 
   // Applies an addition scale to the UV coordinates of the object. Defaults to
   // (1, 1, 1).

src/experimental/platform/hal/renderer.cc

@@ -228,9 +228,13 @@ void Renderer::DoRender(int width, int height) {
       imgui_bridge_->Update();
     }
     if (framebuffer_mode_ == 0) {
-      mjrDrawMode draw_mode = mjDRAW_MODE_COLOR;
+      mjrDrawMode draw_mode = mjDRAW_MODE_DEFAULT;
       if (scene_.flags[mjRND_SEGMENT]) {
-        draw_mode = mjDRAW_MODE_SEGMENTATION;
+        if (scene_.flags[mjRND_IDCOLOR]) {
+          draw_mode = mjDRAW_MODE_SEGMENTATION_BY_ID;
+        } else {
+          draw_mode = mjDRAW_MODE_SEGMENTATION_BY_COLOR;
+        }
       } else if (scene_.flags[mjRND_DEPTH]) {
         draw_mode = mjDRAW_MODE_DEPTH;
       } else if (scene_.flags[mjRND_WIREFRAME]) {
@@ -249,7 +253,7 @@ void Renderer::DoRender(int width, int height) {
 
       mjr_defaultRenderRequest(&reqs[1]);
       reqs[1].scene = imgui_bridge_->GetScene();
-      reqs[1].draw_mode = mjDRAW_MODE_COLOR;
+      reqs[1].draw_mode = mjDRAW_MODE_DEFAULT;
       reqs[1].camera = imgui_bridge_->GetCamera(viewport.width, viewport.height);
       reqs[1].viewport = viewport;
       reqs[1].enable_shadows = false;
@@ -283,9 +287,13 @@ void Renderer::DoReadPixels(int width, int height, unsigned char* rgb) {
     mjr_setBuffer(mjFB_WINDOW, &render_context_);
     FlipImage(rgb, viewport.width, viewport.height, 3);
   } else {
-    mjrDrawMode draw_mode = mjDRAW_MODE_COLOR;
+    mjrDrawMode draw_mode = mjDRAW_MODE_DEFAULT;
     if (scene_.flags[mjRND_SEGMENT]) {
-      draw_mode = mjDRAW_MODE_SEGMENTATION;
+      if (scene_.flags[mjRND_IDCOLOR]) {
+        draw_mode = mjDRAW_MODE_SEGMENTATION_BY_ID;
+      } else {
+        draw_mode = mjDRAW_MODE_SEGMENTATION_BY_COLOR;
+      }
     } else if (scene_.flags[mjRND_DEPTH]) {
       draw_mode = mjDRAW_MODE_DEPTH;
     } else if (scene_.flags[mjRND_WIREFRAME]) {
@@ -312,7 +320,7 @@ void Renderer::DoReadPixels(int width, int height, unsigned char* rgb) {
 
     mjr_defaultRenderRequest(&reqs[1]);
     reqs[1].scene = imgui_bridge_->GetScene();
-    reqs[1].draw_mode = mjDRAW_MODE_COLOR;
+    reqs[1].draw_mode = mjDRAW_MODE_DEFAULT;
     reqs[1].camera = imgui_bridge_->GetCamera(viewport.width, viewport.height);
     reqs[1].target = target.get();
     reqs[1].viewport = viewport;