Add ReadPixels function to RenderTarget.

PiperOrigin-RevId: 900017779
Change-Id: I0632b66124632558ea19157e4ca319e4b0d10617

src/experimental/filament/filament/filament_context.cc

@@ -209,28 +209,6 @@ void FilamentContext::DestroyRenderTargets() {
   color_target_.reset();
 }
 
-static void ReadColorPixels(filament::Renderer* renderer,
-                            RenderTarget* target, mjrRect viewport,
-                            unsigned char* buffer, size_t num_bytes) {
-  filament::backend::PixelBufferDescriptor descriptor(
-      buffer, num_bytes, filament::backend::PixelDataFormat::RGB,
-      filament::backend::PixelDataType::UBYTE);
-  renderer->readPixels(target->GetFilamentRenderTarget(), viewport.left,
-                       viewport.bottom, viewport.width, viewport.height,
-                       std::move(descriptor));
-}
-
-static void ReadDepthPixels(filament::Renderer* renderer,
-                            RenderTarget* target, mjrRect viewport,
-                            float* buffer, size_t num_bytes) {
-  filament::backend::PixelBufferDescriptor descriptor(
-      buffer, num_bytes, filament::backend::PixelDataFormat::R,
-      filament::backend::PixelDataType::FLOAT);
-  renderer->readPixels(target->GetFilamentRenderTarget(), viewport.left,
-                       viewport.bottom, viewport.width, viewport.height,
-                       std::move(descriptor));
-}
-
 void FilamentContext::ReadPixels(mjrRect viewport, unsigned char* rgb,
                                  float* depth) {
   if (scene_swap_chain_target_ != kOffscreenSwapChain) {
@@ -261,7 +239,7 @@ void FilamentContext::ReadPixels(mjrRect viewport, unsigned char* rgb,
       }
 
       const size_t num_bytes = viewport.width * viewport.height * 3;
-      ReadColorPixels(renderer_, color_target_.get(), viewport, rgb, num_bytes);
+      color_target_->ReadColorPixels(renderer_, rgb, num_bytes);
 
       renderer_->endFrame();
     }
@@ -277,8 +255,8 @@ void FilamentContext::ReadPixels(mjrRect viewport, unsigned char* rgb,
       scene_view_->Render(renderer_, request);
 
       const size_t num_bytes = viewport.width * viewport.height * sizeof(float);
-      ReadDepthPixels(renderer_, depth_target_.get(), viewport, depth,
-                      num_bytes);
+      depth_target_->ReadColorPixels(
+          renderer_, reinterpret_cast<uint8_t*>(depth), num_bytes);
 
       renderer_->endFrame();
     }

src/experimental/filament/filament/render_target.cc

@@ -14,18 +14,25 @@
 
 #include "experimental/filament/filament/render_target.h"
 
+#include <cstddef>
+#include <cstdint>
 #include <memory>
+#include <utility>
 
+#include <backend/DriverEnums.h>
+#include <backend/PixelBufferDescriptor.h>
 #include <filament/Engine.h>
+#include <filament/Renderer.h>
 #include <filament/RenderTarget.h>
 #include <filament/Texture.h>
+#include <mujoco/mujoco.h>
 #include "experimental/filament/filament/texture.h"
 
 namespace mujoco {
 
 RenderTarget::RenderTarget(filament::Engine* engine,
-                                                 RenderTargetTextureType color,
-                                                 RenderTargetTextureType depth)
+                           RenderTargetTextureType color,
+                           RenderTargetTextureType depth)
     : engine_(engine), color_type_(color), depth_type_(depth) {}
 
 RenderTarget::~RenderTarget() noexcept {
@@ -53,6 +60,35 @@ void RenderTarget::Prepare(int width, int height) {
   render_target_ = builder.build(*engine_);
 }
 
+void RenderTarget::ReadColorPixels(filament::Renderer* renderer, uint8_t* bytes,
+                                   size_t num_bytes) {
+  filament::backend::PixelDataFormat format;
+  filament::backend::PixelDataType type;
+  size_t expected_num_bytes = 0;
+  switch (color_type_) {
+    case RenderTargetTextureType::kColor:
+      format = filament::backend::PixelDataFormat::RGB;
+      type = filament::backend::PixelDataType::UBYTE;
+      expected_num_bytes = width_ * height_ * 3;
+      break;
+    case RenderTargetTextureType::kDepthColor:
+      format = filament::backend::PixelDataFormat::R;
+      type = filament::backend::PixelDataType::FLOAT;
+      expected_num_bytes = width_ * height_ * sizeof(float);
+      break;
+    default:
+      mju_error("Unsupported pixel format: %d", color_type_);
+      return;
+  }
+  if (num_bytes != expected_num_bytes) {
+    mju_error("Invalid number of bytes.");
+    return;
+  }
+
+  filament::backend::PixelBufferDescriptor desc(bytes, num_bytes, format, type);
+  renderer->readPixels(render_target_, 0, 0, width_, height_, std::move(desc));
+}
+
 void RenderTarget::Destroy() {
   if (render_target_) {
     engine_->destroy(render_target_);

src/experimental/filament/filament/render_target.h

@@ -15,6 +15,8 @@
 #ifndef MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_RENDER_TARGET_H_
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_RENDER_TARGET_H_
 
+#include <cstddef>
+#include <cstdint>
 #include <memory>
 
 #include <filament/Engine.h>
@@ -28,9 +30,8 @@ class RenderTarget {
  public:
   // Defines the types of textures to create for the color and depth
   // attachments.
-  RenderTarget(filament::Engine* engine,
-                          RenderTargetTextureType color,
-                          RenderTargetTextureType depth);
+  RenderTarget(filament::Engine* engine, RenderTargetTextureType color,
+               RenderTargetTextureType depth);
   ~RenderTarget() noexcept;
 
   RenderTarget(const RenderTarget&) = delete;
@@ -40,6 +41,10 @@ class RenderTarget {
   // the last time the render target was prepared.
   void Prepare(int width, int height);
 
+  // Reads the pixels from the render target texture.
+  void ReadColorPixels(filament::Renderer* renderer, uint8_t* bytes,
+                       size_t num_bytes);
+
   // Returns the color texture.
   Texture* GetColorTexture() const;