@@ -33,6 +33,88 @@ bool hasExtension(std::string str, std::string ext) {
3333 }
3434}
3535
36+ // Helper to actually do the render pass and return the result in a buffer
37+ std::vector<unsigned char > getRenderInBuffer (const ScreenshotOptions& options = {}) {
38+ checkInitialized ();
39+
40+ render::engine->useAltDisplayBuffer = true ;
41+ if (options.transparentBackground ) render::engine->lightCopy = true ; // copy directly in to buffer without blending
42+
43+ // == Make sure we render first
44+ processLazyProperties ();
45+
46+ // save the redraw requested bit and restore it below
47+ bool requestedAlready = redrawRequested ();
48+ requestRedraw ();
49+
50+ // Create a new context and push it on to the stack
51+ ImGuiContext* oldContext;
52+ ImGuiContext* newContext;
53+ if (options.includeUI ) {
54+ // WARNING: code duplicated here and in pushContext()
55+ oldContext = ImGui::GetCurrentContext ();
56+ newContext = ImGui::CreateContext ();
57+ ImGuiIO& oldIO = ImGui::GetIO (); // used to GLFW + OpenGL data to the new IO object
58+ #ifdef IMGUI_HAS_DOCK
59+ ImGuiPlatformIO& oldPlatformIO = ImGui::GetPlatformIO ();
60+ #endif
61+ ImGui::SetCurrentContext (newContext);
62+ #ifdef IMGUI_HAS_DOCK
63+ // Propagate GLFW window handle to new context
64+ ImGui::GetMainViewport ()->PlatformHandle = oldPlatformIO.Viewports [0 ]->PlatformHandle ;
65+ #endif
66+ ImGui::GetIO ().BackendPlatformUserData = oldIO.BackendPlatformUserData ;
67+ ImGui::GetIO ().BackendRendererUserData = oldIO.BackendRendererUserData ;
68+
69+ render::engine->configureImGui ();
70+
71+ // render a few times, to let imgui shake itself out
72+ for (int i = 0 ; i < 3 ; i++) {
73+ draw (options.includeUI , false );
74+ }
75+ }
76+
77+ draw (options.includeUI , false );
78+
79+ if (options.includeUI ) {
80+ // WARNING: code duplicated here and in pushContext()
81+ // Workaround overzealous ImGui assertion before destroying any inner context
82+ // https://github.com/ocornut/imgui/pull/7175
83+ ImGui::SetCurrentContext (newContext);
84+ ImGui::GetIO ().BackendPlatformUserData = nullptr ;
85+ ImGui::GetIO ().BackendRendererUserData = nullptr ;
86+
87+ ImGui::DestroyContext (newContext);
88+ ImGui::SetCurrentContext (oldContext);
89+ }
90+
91+
92+ if (requestedAlready) {
93+ requestRedraw ();
94+ }
95+
96+ // these _should_ always be accurate
97+ int w = view::bufferWidth;
98+ int h = view::bufferHeight;
99+ std::vector<unsigned char > buff = render::engine->displayBufferAlt ->readBuffer ();
100+
101+ // Set alpha to 1
102+ if (!options.transparentBackground ) {
103+ for (int j = 0 ; j < h; j++) {
104+ for (int i = 0 ; i < w; i++) {
105+ int ind = i + j * w;
106+ buff[4 * ind + 3 ] = std::numeric_limits<unsigned char >::max ();
107+ }
108+ }
109+ }
110+
111+ render::engine->useAltDisplayBuffer = false ;
112+ if (options.transparentBackground ) render::engine->lightCopy = false ;
113+
114+ return buff;
115+ }
116+
117+
36118} // namespace
37119
38120
@@ -59,111 +141,62 @@ void saveImage(std::string name, unsigned char* buffer, int w, int h, int channe
59141 */
60142
61143 } else {
62- // Fall back on png
63- stbi_write_png (name.c_str (), w, h, channels, buffer, channels * w);
144+ error (" unrecognized file extension, should be one of '.png', '.jpg', '.jpeg'. Got filename: " + name);
64145 }
65146}
66147
67- void screenshot (std::string filename, bool transparentBG ) {
148+ void screenshot (std::string filename, const ScreenshotOptions& options ) {
68149 checkInitialized ();
150+ ScreenshotOptions thisOptions = options; // we may modify it below
69151
70- render::engine->useAltDisplayBuffer = true ;
71- if (transparentBG) render::engine->lightCopy = true ; // copy directly in to buffer without blending
72-
73- // == Make sure we render first
74- processLazyProperties ();
75-
76- // save the redraw requested bit and restore it below
77- bool requestedAlready = redrawRequested ();
78- requestRedraw ();
79-
80- draw (false , false );
81-
82- if (requestedAlready) {
83- requestRedraw ();
152+ // only pngs can be written with transparency
153+ if (!hasExtension (filename, " .png" )) {
154+ thisOptions.transparentBackground = false ;
84155 }
85156
86- // these _should_ always be accurate
157+ std::vector< unsigned char > buff = getRenderInBuffer (thisOptions);
87158 int w = view::bufferWidth;
88159 int h = view::bufferHeight;
89- std::vector<unsigned char > buff = render::engine->displayBufferAlt ->readBuffer ();
90-
91- // Set alpha to 1
92- if (!transparentBG) {
93- for (int j = 0 ; j < h; j++) {
94- for (int i = 0 ; i < w; i++) {
95- int ind = i + j * w;
96- buff[4 * ind + 3 ] = std::numeric_limits<unsigned char >::max ();
97- }
98- }
99- }
100160
101161 // Save to file
102162 saveImage (filename, &(buff.front ()), w, h, 4 );
163+ }
103164
104- render::engine->useAltDisplayBuffer = false ;
105- if (transparentBG) render::engine->lightCopy = false ;
165+ void screenshot (std::string filename, bool transparentBG) {
166+ ScreenshotOptions options;
167+ options.transparentBackground = transparentBG;
168+ screenshot (filename, options);
106169}
107170
108- void screenshot (bool transparentBG ) {
171+ void screenshot (const ScreenshotOptions& options ) {
109172
173+ // construct the filename for the output
110174 char buff[50 ];
111175 snprintf (buff, 50 , " screenshot_%06zu%s" , state::screenshotInd, options::screenshotExtension.c_str ());
112176 std::string defaultName (buff);
113177
114- // only pngs can be written with transparency
115- if (!hasExtension (options::screenshotExtension, " .png" )) {
116- transparentBG = false ;
117- }
118-
119- screenshot (defaultName, transparentBG);
178+ screenshot (defaultName, options);
120179
121180 state::screenshotInd++;
122181}
123182
183+ void screenshot (bool transparentBG) {
184+ ScreenshotOptions options;
185+ options.transparentBackground = transparentBG;
186+ screenshot (options);
187+ }
188+
124189void screenshot (const char * filename) { screenshot (std::string (filename), true ); }
125190
126191void resetScreenshotIndex () { state::screenshotInd = 0 ; }
127192
128193
129- std::vector<unsigned char > screenshotToBuffer (bool transparentBG) {
130- checkInitialized ();
131-
132- render::engine->useAltDisplayBuffer = true ;
133- if (transparentBG) render::engine->lightCopy = true ; // copy directly in to buffer without blending
134-
135- // == Make sure we render first
136- processLazyProperties ();
137-
138- // save the redraw requested bit and restore it below
139- bool requestedAlready = redrawRequested ();
140- requestRedraw ();
141-
142- draw (false , false );
143-
144- if (requestedAlready) {
145- requestRedraw ();
146- }
147-
148- // these _should_ always be accurate
149- int w = view::bufferWidth;
150- int h = view::bufferHeight;
151- std::vector<unsigned char > buff = render::engine->displayBufferAlt ->readBuffer ();
152-
153- // Set alpha to 1
154- if (!transparentBG) {
155- for (int j = 0 ; j < h; j++) {
156- for (int i = 0 ; i < w; i++) {
157- int ind = i + j * w;
158- buff[4 * ind + 3 ] = std::numeric_limits<unsigned char >::max ();
159- }
160- }
161- }
162-
163- render::engine->useAltDisplayBuffer = false ;
164- if (transparentBG) render::engine->lightCopy = false ;
194+ std::vector<unsigned char > screenshotToBuffer (const ScreenshotOptions& options) { return getRenderInBuffer (options); }
165195
166- return buff;
196+ std::vector<unsigned char > screenshotToBuffer (bool transparentBG) {
197+ ScreenshotOptions options;
198+ options.transparentBackground = transparentBG;
199+ return screenshotToBuffer (options);
167200}
168201
169202} // namespace polyscope
0 commit comments