From 17d10d79fc6621a0e99ca78d12e586facfd9e759 Mon Sep 17 00:00:00 2001
From: Mathias Agopian <mathias@google.com>
Date: Tue, 20 Aug 2024 14:07:34 -0700
Subject: [PATCH] new screenshot debug option in gltf_viewer

screenshots are saved in the current directory as
"screenshotXX.ppm" with XX increasing after each
screenshot.
---
 libs/viewer/include/viewer/AutomationEngine.h |  3 +++
 libs/viewer/src/AutomationEngine.cpp          |  5 ++--
 samples/gltf_viewer.cpp                       | 25 ++++++++++++++++---
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/libs/viewer/include/viewer/AutomationEngine.h b/libs/viewer/include/viewer/AutomationEngine.h
index 8747f59d8da0..f0d907a0f7a5 100644
--- a/libs/viewer/include/viewer/AutomationEngine.h
+++ b/libs/viewer/include/viewer/AutomationEngine.h
@@ -224,6 +224,9 @@ class UTILS_PUBLIC AutomationEngine {
      */
     static void exportSettings(const Settings& settings, const char* filename);
 
+    static void exportScreenshot(View* view, Renderer* renderer, std::string filename,
+            bool autoclose, AutomationEngine* automationEngine);
+
     Options getOptions() const { return mOptions; }
     bool isRunning() const { return mIsRunning; }
     size_t currentTest() const { return mCurrentTest; }
diff --git a/libs/viewer/src/AutomationEngine.cpp b/libs/viewer/src/AutomationEngine.cpp
index 19d23ef1681f..810df99dbb8f 100644
--- a/libs/viewer/src/AutomationEngine.cpp
+++ b/libs/viewer/src/AutomationEngine.cpp
@@ -56,7 +56,7 @@ static void convertRGBAtoRGB(void* buffer, uint32_t width, uint32_t height) {
     }
 }
 
-static void exportScreenshot(View* view, Renderer* renderer, std::string filename,
+void AutomationEngine::exportScreenshot(View* view, Renderer* renderer, std::string filename,
         bool autoclose, AutomationEngine* automationEngine) {
     const Viewport& vp = view->getViewport();
     const size_t byteCount = vp.width * vp.height * 4;
@@ -244,7 +244,8 @@ void AutomationEngine::tick(Engine* engine, const ViewerContent& content, float
     }
 
     if (mOptions.exportScreenshots) {
-        exportScreenshot(content.view, content.renderer, prefix + ".ppm", isLastTest, this);
+        AutomationEngine::exportScreenshot(
+                content.view, content.renderer, prefix + ".ppm", isLastTest, this);
     }
 
     if (isLastTest) {
diff --git a/samples/gltf_viewer.cpp b/samples/gltf_viewer.cpp
index b5c8e1b7086c..e3dc4363a383 100644
--- a/samples/gltf_viewer.cpp
+++ b/samples/gltf_viewer.cpp
@@ -61,8 +61,10 @@
 #include <algorithm>
 #include <array>
 #include <fstream>
+#include <iomanip>
 #include <iostream>
 #include <set>
+#include <sstream>
 #include <string>
 
 #include "generated/resources/gltf_demo.h"
@@ -137,6 +139,8 @@ struct App {
 
     AutomationSpec* automationSpec = nullptr;
     AutomationEngine* automationEngine = nullptr;
+    bool screenshot = false;
+    uint8_t screenshotSeq = 0;
 };
 
 static const char* DEFAULT_IBL = "assets/ibl/lightroom_14b";
@@ -877,10 +881,15 @@ int main(int argc, char** argv) {
 
             if (ImGui::CollapsingHeader("Debug")) {
                 auto& debug = engine->getDebugRegistry();
-                if (ImGui::Button("Capture frame")) {
-                    bool* captureFrame =
-                            debug.getPropertyAddress<bool>("d.renderer.doFrameCapture");
-                    *captureFrame = true;
+                if (engine->getBackend() == Engine::Backend::METAL) {
+                    if (ImGui::Button("Capture frame")) {
+                        bool* captureFrame =
+                                debug.getPropertyAddress<bool>("d.renderer.doFrameCapture");
+                        *captureFrame = true;
+                    }
+                }
+                if (ImGui::Button("Screenshot")) {
+                    app.screenshot = true;
                 }
                 ImGui::Checkbox("Disable buffer padding",
                         debug.getPropertyAddress<bool>("d.renderer.disable_buffer_padding"));
@@ -1138,6 +1147,14 @@ int main(int argc, char** argv) {
     };
 
     auto postRender = [&app](Engine* engine, View* view, Scene*, Renderer* renderer) {
+        if (app.screenshot) {
+            std::ostringstream stringStream;
+            stringStream << "screenshot" << std::setfill('0') << std::setw(2) << +app.screenshotSeq;
+            AutomationEngine::exportScreenshot(
+                    view, renderer, stringStream.str() + ".ppm", false, app.automationEngine);
+            ++app.screenshotSeq;
+            app.screenshot = false;
+        }
         if (app.automationEngine->shouldClose()) {
             FilamentApp::get().close();
             return;