From 8507b62556aa2e53400907465823c7774388e6c7 Mon Sep 17 00:00:00 2001
From: Sam Haves <shaves@google.com>
Date: Wed, 3 Dec 2025 16:05:34 -0800
Subject: [PATCH] Cache the rendering mode used to allow headless segmentation
 rendering.

This also moves the read call into the frame, per the recommendation in the filament documentation. We also don't need to flush and wait for the engine twice, just defer to the end of readPixels.

PiperOrigin-RevId: 839953612
Change-Id: I52205147794440322ca976c52fb4905207c21473
---
 .../filament/filament/filament_context.cc     | 30 +++++++++----------
 .../filament/filament/filament_context.h      |  1 +
 2 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/src/experimental/filament/filament/filament_context.cc b/src/experimental/filament/filament/filament_context.cc
index 0623723b..e06e8b42 100644
--- a/src/experimental/filament/filament/filament_context.cc
+++ b/src/experimental/filament/filament/filament_context.cc
@@ -151,13 +151,12 @@ void FilamentContext::Render(const mjrRect& viewport, const mjvScene* scene,
     render_gui_ = gui_view_->PrepareRenderable();
   }
 
-  // Render the frame if we're not rendering to a texture.s
+  last_render_mode_ = scene->flags[mjRND_SEGMENT]
+                                 ? SceneView::DrawMode::kSegmentation
+                                 : SceneView::DrawMode::kNormal;
+  // Render the frame if we're not rendering to a texture.
   if (!render_to_texture_) {
-    SceneView::DrawMode mode = scene->flags[mjRND_SEGMENT]
-                                   ? SceneView::DrawMode::kSegmentation
-                                   : SceneView::DrawMode::kNormal;
-
-    filament::View* view = scene_view_->PrepareRenderView(mode);
+    filament::View* view = scene_view_->PrepareRenderView(last_render_mode_);
 
     // Wait until previous frame is completed before requesting a new frame.
     engine_->flushAndWait();
@@ -260,7 +259,7 @@ void FilamentContext::ReadPixels(mjrRect viewport, unsigned char* rgb,
 
   if (rgb) {
     filament::View* view =
-        scene_view_->PrepareRenderView(SceneView::DrawMode::kNormal);
+        scene_view_->PrepareRenderView(last_render_mode_);
     if (renderer_->beginFrame(swap_chain_)) {
       // We need to disable msaa in order to render to texture.
       auto options = view->getMultiSampleAntiAliasingOptions();
@@ -269,14 +268,14 @@ void FilamentContext::ReadPixels(mjrRect viewport, unsigned char* rgb,
       });
       view->setRenderTarget(color_target_);
       renderer_->render(view);
+
+      const size_t num_bytes = viewport.width * viewport.height * 3;
+      ReadColorPixels(renderer_, color_target_, viewport, rgb, num_bytes);
+
       view->setRenderTarget(nullptr);
       view->setMultiSampleAntiAliasingOptions(options);
       renderer_->endFrame();
     }
-
-    engine_->flushAndWait();
-    const size_t num_bytes = viewport.width * viewport.height * 3;
-    ReadColorPixels(renderer_, color_target_, viewport, rgb, num_bytes);
   }
 
   if (depth) {
@@ -285,16 +284,17 @@ void FilamentContext::ReadPixels(mjrRect viewport, unsigned char* rgb,
     if (renderer_->beginFrame(swap_chain_)) {
       view->setRenderTarget(depth_target_);
       renderer_->render(view);
+
+      const size_t num_bytes = viewport.width * viewport.height * sizeof(float);
+      ReadDepthPixels(renderer_, depth_target_, viewport, depth, num_bytes);
+
       view->setRenderTarget(nullptr);
       renderer_->endFrame();
     }
-
-    engine_->flushAndWait();
-    const size_t num_bytes = viewport.width * viewport.height * sizeof(float);
-    ReadDepthPixels(renderer_, depth_target_, viewport, depth, num_bytes);
   }
 
   if (rgb || depth) {
+    // Wait for rendering and copy back to buffer to complete.
     engine_->flushAndWait();
   }
 }
diff --git a/src/experimental/filament/filament/filament_context.h b/src/experimental/filament/filament/filament_context.h
index 5aa6e28d..b5829a3d 100644
--- a/src/experimental/filament/filament/filament_context.h
+++ b/src/experimental/filament/filament/filament_context.h
@@ -75,6 +75,7 @@ class FilamentContext {
   filament::Texture* target_textures_[kNumRenderTargetTextureTypes] = {
       nullptr, nullptr, nullptr};
 
+  SceneView::DrawMode last_render_mode_ = SceneView::DrawMode::kNormal;
   bool render_to_texture_ = false;
   bool render_gui_ = false;