From 2a6eda72f1b79b277dfe28ebc667e7bd17a54caf Mon Sep 17 00:00:00 2001
From: Nicholas Sharp <nmwsharp@gmail.com>
Date: Wed, 27 Mar 2024 22:59:55 -0700
Subject: [PATCH] add screenshot to buffer

---
 include/polyscope/screenshot.h |  9 +++++---
 src/screenshot.cpp             | 40 ++++++++++++++++++++++++++++++++++
 test/src/basics_test.cpp       | 10 +++++++--
 3 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/include/polyscope/screenshot.h b/include/polyscope/screenshot.h
index 27c3f097..647c85b6 100644
--- a/include/polyscope/screenshot.h
+++ b/include/polyscope/screenshot.h
@@ -7,12 +7,15 @@
 namespace polyscope {
 
 
-// Take screenshots of the current view
+// Take a screenshot from the current view and write to file
 void screenshot(std::string filename, bool transparentBG = true);
-void screenshot(bool transparentBG = true);
-void saveImage(std::string name, unsigned char* buffer, int w, int h, int channels);
+void screenshot(bool transparentBG = true); // automatic file names like `screenshot_000000.png`
+void saveImage(std::string name, unsigned char* buffer, int w, int h, int channels); // helper
 void resetScreenshotIndex();
 
+// Take a screenshot from the current view and return it as a buffer
+// the dimensions are view::bufferWidth and view::bufferHeight , with entries RGBA at 1 byte each.
+std::vector<unsigned char> screenshotToBuffer(bool transparentBG = true);
 
 namespace state {
 
diff --git a/src/screenshot.cpp b/src/screenshot.cpp
index a9a4a0d2..6015eb1e 100644
--- a/src/screenshot.cpp
+++ b/src/screenshot.cpp
@@ -121,4 +121,44 @@ void screenshot(bool transparentBG) {
 
 void resetScreenshotIndex() { state::screenshotInd = 0; }
 
+
+std::vector<unsigned char> screenshotToBuffer(bool transparentBG) {
+
+  render::engine->useAltDisplayBuffer = true;
+  if (transparentBG) render::engine->lightCopy = true; // copy directly in to buffer without blending
+
+  // == Make sure we render first
+  processLazyProperties();
+
+  // save the redraw requested bit and restore it below
+  bool requestedAlready = redrawRequested();
+  requestRedraw();
+
+  draw(false, false);
+
+  if (requestedAlready) {
+    requestRedraw();
+  }
+
+  // these _should_ always be accurate
+  int w = view::bufferWidth;
+  int h = view::bufferHeight;
+  std::vector<unsigned char> buff = render::engine->displayBufferAlt->readBuffer();
+
+  // Set alpha to 1
+  if (!transparentBG) {
+    for (int j = 0; j < h; j++) {
+      for (int i = 0; i < w; i++) {
+        int ind = i + j * w;
+        buff[4 * ind + 3] = std::numeric_limits<unsigned char>::max();
+      }
+    }
+  }
+
+  render::engine->useAltDisplayBuffer = false;
+  if (transparentBG) render::engine->lightCopy = false;
+
+  return buff;
+}
+
 } // namespace polyscope
diff --git a/test/src/basics_test.cpp b/test/src/basics_test.cpp
index 9c1efb58..a8248a5f 100644
--- a/test/src/basics_test.cpp
+++ b/test/src/basics_test.cpp
@@ -99,8 +99,14 @@ TEST_F(PolyscopeTest, EmptyBuffer) {
   polyscope::removeAllStructures();
 }
 
-TEST_F(PolyscopeTest, Screenshot) {
-  polyscope::screenshot("test_screeshot.png");
+TEST_F(PolyscopeTest, Screenshot) { polyscope::screenshot("test_screeshot.png"); }
+
+TEST_F(PolyscopeTest, ScreenshotBuffer) {
+  std::vector<unsigned char> buff = polyscope::screenshotToBuffer();
+  EXPECT_EQ(buff.size(), polyscope::view::bufferWidth * polyscope::view::bufferHeight * 4);
+
+  std::vector<unsigned char> buff2 = polyscope::screenshotToBuffer(false);
+  EXPECT_EQ(buff2.size(), polyscope::view::bufferWidth * polyscope::view::bufferHeight * 4);
 }
 
 // ============================================================