add screenshot to buffer

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 {
 

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

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);
 }
 
 // ============================================================