chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Move segmentation coloring logic into renderer.

Browse Source 
PiperOrigin-RevId: 927300828
Change-Id: I994c5742ec30993d2245df167042e72820f91ac7
This commit is contained in:
Haroon Qureshi
2026-06-05 07:57:28 -07:00
committed by Copybara-Service
parent fdd82ac420
commit 511cf27ff8
8 changed files with 52 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/experimental/filament/compat/scene_geom_util.cc
+1 -14
View File
@@ -183,20 +183,7 @@ static void UpdateGeomMaterial(mjrRenderable* renderable, const mjvGeom& geom,
material.roughness = model->mat_roughness[geom.matid];
}
if (geom.segid >= 0) {
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);
}
material.segmentation_id = geom.segid;
// 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
+9 -5
View File
@@ -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{
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};
const uint8_t r = material.segmentation_id >> 0;
const uint8_t g = material.segmentation_id >> 8;
const uint8_t b = material.segmentation_id >> 16;
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
+10 -5
View File
@@ -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_color[1] = 0;
material.segmentation_color[2] = 0;
material.segmentation_id = 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
+2 -1
View File
@@ -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
+7 -3
View File
@@ -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
-3
View File
@@ -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
+9 -7
View File
@@ -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.
mjDRAW_MODE_COLOR,
// Render the scene with default settings for colors and lighting (shading).
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
// lighting and texture.
mjDRAW_MODE_SEGMENTATION,
// Render each object using its segmentation id (see mjrMaterial).
mjDRAW_MODE_SEGMENTATION_BY_ID,
// 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.
mjtByte segmentation_color[3];
// The ID to use for segmentation rendering. These IDs will be mapped to a
// 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
+14 -6
View File
@@ -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;