Skip to content

Commit d7a2a04

Browse files
committed
clean up screenshots API, add support for including UI
1 parent 258cd7b commit d7a2a04

File tree

6 files changed

+145
-81
lines changed

6 files changed

+145
-81
lines changed

include/polyscope/options.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,8 @@ extern float shadowDarkness;
8585

8686
extern bool screenshotTransparency; // controls whether screenshots taken by clicking the GUI button have a
8787
// transparent background
88+
extern bool screenshotWithImGuiUI; // controls whether screenshots taken by clicking the GUI button have a
89+
// transparent background
8890
extern std::string screenshotExtension; // sets the extension used for automatically-numbered screenshots (e.g. by
8991
// clicking the GUI button)
9092

include/polyscope/screenshot.h

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,17 +6,30 @@
66

77
namespace polyscope {
88

9+
struct ScreenshotOptions {
10+
bool transparentBackground = true;
11+
bool includeUI = false;
12+
};
913

10-
// Take a screenshot from the current view and write to file
11-
void screenshot(bool transparentBG = true); // automatic file names like `screenshot_000000.png`
14+
// Save a screenshot to a file
15+
void screenshot(const ScreenshotOptions& options = {}); // automatic file names like `screenshot_000000.png`
16+
void screenshot(std::string filename, const ScreenshotOptions& options = {});
17+
18+
// Save a screenshot to a buffer
19+
std::vector<unsigned char> screenshotToBuffer(const ScreenshotOptions& options = {});
20+
21+
22+
// (below: various legacy versions of the function, prefer the general form above))
23+
24+
void screenshot(bool transparentBG); // automatic file names like `screenshot_000000.png`
1225
void screenshot(std::string filename, bool transparentBG = true);
1326
void screenshot(const char* filename); // this is needed because annoyingly overload resolution prefers the bool version
1427
void saveImage(std::string name, unsigned char* buffer, int w, int h, int channels); // helper
1528
void resetScreenshotIndex();
1629

1730
// Take a screenshot from the current view and return it as a buffer
1831
// the dimensions are view::bufferWidth and view::bufferHeight , with entries RGBA at 1 byte each.
19-
std::vector<unsigned char> screenshotToBuffer(bool transparentBG = true);
32+
std::vector<unsigned char> screenshotToBuffer(bool transparentBG);
2033

2134
namespace state {
2235

src/options.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ bool warnForInvalidValues = true;
3434
bool displayMessagePopups = true;
3535

3636
bool screenshotTransparency = true;
37+
bool screenshotWithImGuiUI = false;
3738
std::string screenshotExtension = ".png";
3839

3940
// == Scene options

src/polyscope.cpp

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -201,6 +201,8 @@ bool isInitialized() { return state::initialized; }
201201

202202
void pushContext(std::function<void()> callbackFunction, bool drawDefaultUI) {
203203

204+
// WARNING: code duplicated here and in screenshot.cpp
205+
204206
// Create a new context and push it on to the stack
205207
ImGuiContext* newContext = ImGui::CreateContext();
206208
ImGuiIO& oldIO = ImGui::GetIO(); // used to GLFW + OpenGL data to the new IO object
@@ -222,7 +224,7 @@ void pushContext(std::function<void()> callbackFunction, bool drawDefaultUI) {
222224
if (contextStack.size() > 50) {
223225
// Catch bugs with nested show()
224226
exception("Uh oh, polyscope::show() was recusively MANY times (depth > 50), this is probably a bug. Perhaps "
225-
"you are accidentally calling show every time polyscope::userCallback executes?");
227+
"you are accidentally calling show() every time polyscope::userCallback executes?");
226228
};
227229

228230
// Make sure the window is visible
@@ -253,6 +255,7 @@ void pushContext(std::function<void()> callbackFunction, bool drawDefaultUI) {
253255
}
254256
}
255257

258+
// WARNING: code duplicated here and in screenshot.cpp
256259
// Workaround overzealous ImGui assertion before destroying any inner context
257260
// https://github.com/ocornut/imgui/pull/7175
258261
ImGui::SetCurrentContext(newContext);
@@ -621,7 +624,10 @@ void buildPolyscopeGui() {
621624
ImGui::SameLine();
622625
ImGui::PushStyleVar(ImGuiStyleVar_ItemSpacing, ImVec2(1.0f, 0.0f));
623626
if (ImGui::Button("Screenshot")) {
624-
screenshot(options::screenshotTransparency);
627+
ScreenshotOptions options;
628+
options.transparentBackground = options::screenshotTransparency;
629+
options.includeUI = options::screenshotWithImGuiUI;
630+
screenshot(options);
625631
}
626632
ImGui::SameLine();
627633
if (ImGui::ArrowButton("##Option", ImGuiDir_Down)) {
@@ -631,6 +637,7 @@ void buildPolyscopeGui() {
631637
if (ImGui::BeginPopup("ScreenshotOptionsPopup")) {
632638

633639
ImGui::Checkbox("with transparency", &options::screenshotTransparency);
640+
ImGui::Checkbox("with UI", &options::screenshotWithImGuiUI);
634641

635642
if (ImGui::BeginMenu("file format")) {
636643
if (ImGui::MenuItem(".png", NULL, options::screenshotExtension == ".png")) options::screenshotExtension = ".png";

src/screenshot.cpp

Lines changed: 108 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -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+
124189
void screenshot(const char* filename) { screenshot(std::string(filename), true); }
125190

126191
void 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

test/src/basics_test.cpp

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,15 @@ TEST_F(PolyscopeTest, WindowProperties) {
130130
polyscope::show(3);
131131
}
132132

133-
TEST_F(PolyscopeTest, Screenshot) { polyscope::screenshot("test_screeshot.png"); }
133+
TEST_F(PolyscopeTest, Screenshot) {
134+
polyscope::screenshot("test_screeshot.png");
135+
polyscope::screenshot();
136+
137+
polyscope::ScreenshotOptions opts;
138+
opts.includeUI = true;
139+
opts.transparentBackground = false;
140+
polyscope::screenshot(opts);
141+
}
134142

135143
TEST_F(PolyscopeTest, ScreenshotBuffer) {
136144
std::vector<unsigned char> buff = polyscope::screenshotToBuffer();

0 commit comments

Comments
 (0)