From d5872acec8d7405f95862e2d199950f59147ccd8 Mon Sep 17 00:00:00 2001
From: Haroon Qureshi <haroonq@google.com>
Date: Fri, 15 May 2026 10:08:35 -0700
Subject: [PATCH] Use RGB8 for segmentation colors.

PiperOrigin-RevId: 916061579
Change-Id: I0c1270cb71a9a4476a54b3bf6049981ed30f229f
---
 src/experimental/filament/compat/scene_geom_util.cc  | 10 +++-------
 src/experimental/filament/filament/material.cc       |  7 ++++++-
 src/experimental/filament/render_context_filament.cc |  4 +++-
 src/experimental/filament/render_context_filament.h  |  2 +-
 4 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/src/experimental/filament/compat/scene_geom_util.cc b/src/experimental/filament/compat/scene_geom_util.cc
index 0aa2c56a..0dc9d1f0 100644
--- a/src/experimental/filament/compat/scene_geom_util.cc
+++ b/src/experimental/filament/compat/scene_geom_util.cc
@@ -226,13 +226,9 @@ static void UpdateGeomMaterial(mjrRenderable* renderable, const mjvGeom& geom,
       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 >> 16) & 0xff;
-    material.segmentation_color[0] = static_cast<float>(red) / 255.0f;
-    material.segmentation_color[1] = static_cast<float>(green) / 255.0f;
-    material.segmentation_color[2] = static_cast<float>(blue) / 255.0f;
+    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
diff --git a/src/experimental/filament/filament/material.cc b/src/experimental/filament/filament/material.cc
index aded053b..b96f6f17 100644
--- a/src/experimental/filament/filament/material.cc
+++ b/src/experimental/filament/filament/material.cc
@@ -19,6 +19,7 @@
 #include <filament/MaterialInstance.h>
 #include <filament/RenderableManager.h>
 #include <filament/TextureSampler.h>
+#include <math/vec4.h>
 #include <mujoco/mujoco.h>
 #include "experimental/filament/filament_util.h"
 #include "experimental/filament/filament/mesh.h"
@@ -105,8 +106,12 @@ void UpdateMaterialInstance(filament::MaterialInstance* instance,
                            ReadFloat4(material.color));
   }
   if (fmaterial->hasParameter("SegmentationColor")) {
+    filament::math::float4 color{
+        static_cast<float>(material.segmentation_color[0]) / 255.0f,
+        static_cast<float>(material.segmentation_color[1]) / 255.0f,
+        static_cast<float>(material.segmentation_color[2]) / 255.0f, 1.0f};
     instance->setParameter("SegmentationColor", filament::RgbaType::LINEAR,
-                           ReadFloat4(material.segmentation_color));
+                           color);
   }
   if (fmaterial->hasParameter("EmissiveFactor")) {
     instance->setParameter("EmissiveFactor", material.emissive);
diff --git a/src/experimental/filament/render_context_filament.cc b/src/experimental/filament/render_context_filament.cc
index c2d10cf3..c6b66a7a 100644
--- a/src/experimental/filament/render_context_filament.cc
+++ b/src/experimental/filament/render_context_filament.cc
@@ -97,8 +97,10 @@ void mjr_defaultLightParams(mjrLightParams* params) {
 void mjr_defaultMaterial(mjrMaterial* material) {
   memset(material, 0, sizeof(mjrMaterial));
   setf(material->color, {1.f, 1.f, 1.f, 1.f});
-  setf(material->segmentation_color, {1, 1, 1, 1});
   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;
diff --git a/src/experimental/filament/render_context_filament.h b/src/experimental/filament/render_context_filament.h
index 94c49e6b..43c7fb88 100644
--- a/src/experimental/filament/render_context_filament.h
+++ b/src/experimental/filament/render_context_filament.h
@@ -590,7 +590,7 @@ struct mjrMaterial {
   float color[4];
 
   // The color to use for segmentation rendering. Defaults to white.
-  float segmentation_color[4];
+  mjtByte segmentation_color[3];
 
   // Applies an addition scale to the UV coordinates of the object. Defaults to
   // (1, 1, 1).