UI: Add ability to make screenshots

Co-authored-by: Jim <obs.jim@gmail.com>

UI/CMakeLists.txt

@@ -200,6 +200,7 @@ set(obs_SOURCES
 	window-basic-main-outputs.cpp
 	window-basic-source-select.cpp
 	window-basic-settings-stream.cpp
+	window-basic-main-screenshot.cpp
 	window-basic-auto-config-test.cpp
 	window-basic-main-scene-collections.cpp
 	window-basic-main-transitions.cpp
@@ -288,6 +289,7 @@ set(obs_HEADERS
 	mute-checkbox.hpp
 	record-button.hpp
 	ui-validation.hpp
+	screenshot-obj.hpp
 	url-push-button.hpp
 	volume-control.hpp
 	adv-audio-control.hpp

UI/api-interface.cpp

@@ -534,6 +534,17 @@ struct OBSStudioAPI : obs_frontend_callbacks {
 		}
 	}
 
+	void obs_frontend_take_screenshot(void) override
+	{
+		QMetaObject::invokeMethod(main, "Screenshot");
+	}
+
+	void obs_frontend_take_source_screenshot(obs_source_t *source) override
+	{
+		QMetaObject::invokeMethod(main, "Screenshot",
+					  Q_ARG(OBSSource, OBSSource(source)));
+	}
+
 	void on_load(obs_data_t *settings) override
 	{
 		for (size_t i = saveCallbacks.size(); i > 0; i--) {

UI/data/locale/en-US.ini

@@ -787,6 +787,14 @@ Basic.Settings.Output.Adv.FFmpeg.MuxerSettings="Muxer Settings (if any)"
 Basic.Settings.Output.Adv.FFmpeg.GOPSize="Keyframe interval (frames)"
 Basic.Settings.Output.Adv.FFmpeg.IgnoreCodecCompat="Show all codecs (even if potentially incompatible)"
 
+# Screenshot
+Screenshot="Screenshot Output"
+Screenshot.SourceHotkey="Screenshot Selected Source"
+Screenshot.StudioProgram="Screenshot (Program)"
+Screenshot.Preview="Screenshot (Preview)"
+Screenshot.Scene="Screenshot (Scene)"
+Screenshot.Source="Screenshot (Source)"
+
 # basic mode 'output' settings - advanced section - recording subsection - completer
 FilenameFormatting.completer="%CCYY-%MM-%DD %hh-%mm-%ss\n%YY-%MM-%DD %hh-%mm-%ss\n%Y-%m-%d %H-%M-%S\n%y-%m-%d %H-%M-%S\n%a %Y-%m-%d %H-%M-%S\n%A %Y-%m-%d %H-%M-%S\n%Y-%b-%d %H-%M-%S\n%Y-%B-%d %H-%M-%S\n%Y-%m-%d %I-%M-%S-%p\n%Y-%m-%d %H-%M-%S-%z\n%Y-%m-%d %H-%M-%S-%Z\n%FPS\n%CRES\n%ORES\n%VF"
 

UI/obs-frontend-api/obs-frontend-api.cpp

@@ -463,3 +463,15 @@ void obs_frontend_set_current_preview_scene(obs_source_t *scene)
 	if (callbacks_valid())
 		c->obs_frontend_set_current_preview_scene(scene);
 }
+
+void obs_frontend_take_screenshot(void)
+{
+	if (callbacks_valid())
+		c->obs_frontend_take_screenshot();
+}
+
+void obs_frontend_take_source_screenshot(obs_source_t *source)
+{
+	if (callbacks_valid())
+		c->obs_frontend_take_source_screenshot(source);
+}

UI/obs-frontend-api/obs-frontend-api.h

@@ -194,6 +194,9 @@ EXPORT bool obs_frontend_preview_enabled(void);
 EXPORT obs_source_t *obs_frontend_get_current_preview_scene(void);
 EXPORT void obs_frontend_set_current_preview_scene(obs_source_t *scene);
 
+EXPORT void obs_frontend_take_screenshot(void);
+EXPORT void obs_frontend_take_source_screenshot(obs_source_t *source);
+
 /* ------------------------------------------------------------------------- */
 
 #ifdef __cplusplus

UI/obs-frontend-api/obs-frontend-internal.hpp

@@ -118,6 +118,10 @@ struct obs_frontend_callbacks {
 	virtual void on_preload(obs_data_t *settings) = 0;
 	virtual void on_save(obs_data_t *settings) = 0;
 	virtual void on_event(enum obs_frontend_event event) = 0;
+
+	virtual void obs_frontend_take_screenshot() = 0;
+	virtual void
+	obs_frontend_take_source_screenshot(obs_source_t *source) = 0;
 };
 
 EXPORT void

UI/screenshot-obj.hpp

@@ -0,0 +1,47 @@
+/******************************************************************************
+    Copyright (C) 2020 by Hugh Bailey <obs.jim@gmail.com>
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 2 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+******************************************************************************/
+
+#include <QObject>
+#include <string>
+#include <thread>
+#include <obs.hpp>
+
+class ScreenshotObj : public QObject {
+	Q_OBJECT
+
+public:
+	ScreenshotObj(obs_source_t *source);
+	~ScreenshotObj() override;
+	void Screenshot();
+	void Download();
+	void Copy();
+	void MuxAndFinish();
+
+	gs_texrender_t *texrender = nullptr;
+	gs_stagesurf_t *stagesurf = nullptr;
+	OBSWeakSource weakSource;
+	std::string path;
+	QImage image;
+	uint32_t cx;
+	uint32_t cy;
+	std::thread th;
+
+	int stage = 0;
+
+public slots:
+	void Save();
+};

UI/window-basic-main-screenshot.cpp

@@ -0,0 +1,217 @@
+/******************************************************************************
+    Copyright (C) 2020 by Hugh Bailey <obs.jim@gmail.com>
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 2 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+******************************************************************************/
+
+#include "window-basic-main.hpp"
+#include "screenshot-obj.hpp"
+#include "qt-wrappers.hpp"
+
+static void ScreenshotTick(void *param, float);
+
+/* ========================================================================= */
+
+ScreenshotObj::ScreenshotObj(obs_source_t *source)
+	: weakSource(OBSGetWeakRef(source))
+{
+	obs_add_tick_callback(ScreenshotTick, this);
+}
+
+ScreenshotObj::~ScreenshotObj()
+{
+	obs_enter_graphics();
+	gs_stagesurface_destroy(stagesurf);
+	gs_texrender_destroy(texrender);
+	obs_leave_graphics();
+
+	obs_remove_tick_callback(ScreenshotTick, this);
+	if (th.joinable())
+		th.join();
+}
+
+void ScreenshotObj::Screenshot()
+{
+	OBSSource source = OBSGetStrongRef(weakSource);
+
+	if (source) {
+		cx = obs_source_get_base_width(source);
+		cy = obs_source_get_base_height(source);
+	} else {
+		obs_video_info ovi;
+		obs_get_video_info(&ovi);
+		cx = ovi.base_width;
+		cy = ovi.base_height;
+	}
+
+	if (!cx || !cy) {
+		blog(LOG_WARNING, "Cannot screenshot, invalid target size");
+		obs_remove_tick_callback(ScreenshotTick, this);
+		deleteLater();
+		return;
+	}
+
+	texrender = gs_texrender_create(GS_RGBA, GS_ZS_NONE);
+	stagesurf = gs_stagesurface_create(cx, cy, GS_RGBA);
+
+	gs_texrender_reset(texrender);
+	if (gs_texrender_begin(texrender, cx, cy)) {
+		vec4 zero;
+		vec4_zero(&zero);
+
+		gs_clear(GS_CLEAR_COLOR, &zero, 0.0f, 0);
+		gs_ortho(0.0f, (float)cx, 0.0f, (float)cy, -100.0f, 100.0f);
+
+		gs_blend_state_push();
+		gs_blend_function(GS_BLEND_ONE, GS_BLEND_ZERO);
+
+		if (source) {
+			obs_source_inc_showing(source);
+			obs_source_video_render(source);
+			obs_source_dec_showing(source);
+		} else {
+			obs_render_main_texture();
+		}
+
+		gs_blend_state_pop();
+		gs_texrender_end(texrender);
+	}
+}
+
+void ScreenshotObj::Download()
+{
+	gs_stage_texture(stagesurf, gs_texrender_get_texture(texrender));
+}
+
+void ScreenshotObj::Copy()
+{
+	uint8_t *videoData = nullptr;
+	uint32_t videoLinesize = 0;
+	bool success = false;
+
+	image = QImage(cx, cy, QImage::Format::Format_RGBX8888);
+
+	if (gs_stagesurface_map(stagesurf, &videoData, &videoLinesize)) {
+		int linesize = image.bytesPerLine();
+		for (int y = 0; y < (int)cy; y++)
+			memcpy(image.scanLine(y),
+			       videoData + (y * videoLinesize), linesize);
+
+		gs_stagesurface_unmap(stagesurf);
+		success = true;
+	}
+}
+
+void ScreenshotObj::Save()
+{
+	OBSBasic *main = OBSBasic::Get();
+	config_t *config = main->Config();
+
+	const char *mode = config_get_string(config, "Output", "Mode");
+	const char *type = config_get_string(config, "AdvOut", "RecType");
+	const char *adv_path =
+		strcmp(type, "Standard")
+			? config_get_string(config, "AdvOut", "FFFilePath")
+			: config_get_string(config, "AdvOut", "RecFilePath");
+	const char *rec_path =
+		strcmp(mode, "Advanced")
+			? config_get_string(config, "SimpleOutput", "FilePath")
+			: adv_path;
+
+	const char *filenameFormat =
+		config_get_string(config, "Output", "FilenameFormatting");
+	bool overwriteIfExists =
+		config_get_bool(config, "Output", "OverwriteIfExists");
+
+	path = GetOutputFilename(
+		rec_path, "png", false, overwriteIfExists,
+		GetFormatString(filenameFormat, "Screenshot", nullptr).c_str());
+
+	th = std::thread([this] { MuxAndFinish(); });
+}
+
+void ScreenshotObj::MuxAndFinish()
+{
+	image.save(QT_UTF8(path.c_str()));
+	blog(LOG_INFO, "Saved screenshot to '%s'", path.c_str());
+	deleteLater();
+}
+
+/* ========================================================================= */
+
+#define STAGE_SCREENSHOT 0
+#define STAGE_DOWNLOAD 1
+#define STAGE_COPY_AND_SAVE 2
+#define STAGE_FINISH 3
+
+static void ScreenshotTick(void *param, float)
+{
+	ScreenshotObj *data = reinterpret_cast<ScreenshotObj *>(param);
+
+	if (data->stage == STAGE_FINISH) {
+		return;
+	}
+
+	obs_enter_graphics();
+
+	switch (data->stage) {
+	case STAGE_SCREENSHOT:
+		data->Screenshot();
+		break;
+	case STAGE_DOWNLOAD:
+		data->Download();
+		break;
+	case STAGE_COPY_AND_SAVE:
+		data->Copy();
+		QMetaObject::invokeMethod(data, "Save");
+		obs_remove_tick_callback(ScreenshotTick, data);
+		break;
+	}
+
+	obs_leave_graphics();
+
+	data->stage++;
+}
+
+void OBSBasic::Screenshot(OBSSource source)
+{
+	if (!!screenshotData) {
+		blog(LOG_WARNING, "Cannot take new screenshot, "
+				  "screenshot currently in progress");
+		return;
+	}
+
+	screenshotData = new ScreenshotObj(source);
+}
+
+void OBSBasic::ScreenshotSelectedSource()
+{
+	OBSSceneItem item = GetCurrentSceneItem();
+	if (item) {
+		Screenshot(obs_sceneitem_get_source(item));
+	} else {
+		blog(LOG_INFO, "Could not take a source screenshot: "
+			       "no source selected");
+	}
+}
+
+void OBSBasic::ScreenshotProgram()
+{
+	Screenshot(GetProgramSource());
+}
+
+void OBSBasic::ScreenshotScene()
+{
+	Screenshot(GetCurrentSceneSource());
+}

UI/window-basic-main.cpp

@@ -2347,6 +2347,31 @@ void OBSBasic::CreateHotkeys()
 		"OBSBasic.ResetStats", Str("Basic.Stats.ResetStats"),
 		resetStats, this);
 	LoadHotkey(statsHotkey, "OBSBasic.ResetStats");
+
+	auto screenshot = [](void *data, obs_hotkey_id, obs_hotkey_t *,
+			     bool pressed) {
+		if (pressed)
+			QMetaObject::invokeMethod(static_cast<OBSBasic *>(data),
+						  "Screenshot",
+						  Qt::QueuedConnection);
+	};
+
+	screenshotHotkey = obs_hotkey_register_frontend(
+		"OBSBasic.Screenshot", Str("Screenshot"), screenshot, this);
+	LoadHotkey(screenshotHotkey, "OBSBasic.Screenshot");
+
+	auto screenshotSource = [](void *data, obs_hotkey_id, obs_hotkey_t *,
+				   bool pressed) {
+		if (pressed)
+			QMetaObject::invokeMethod(static_cast<OBSBasic *>(data),
+						  "ScreenshotSelectedSource",
+						  Qt::QueuedConnection);
+	};
+
+	sourceScreenshotHotkey = obs_hotkey_register_frontend(
+		"OBSBasic.SelectedSourceScreenshot",
+		Str("Screenshot.SourceHotkey"), screenshotSource, this);
+	LoadHotkey(sourceScreenshotHotkey, "OBSBasic.SelectedSourceScreenshot");
 }
 
 void OBSBasic::ClearHotkeys()
@@ -2360,6 +2385,8 @@ void OBSBasic::ClearHotkeys()
 	obs_hotkey_unregister(togglePreviewProgramHotkey);
 	obs_hotkey_unregister(transitionHotkey);
 	obs_hotkey_unregister(statsHotkey);
+	obs_hotkey_unregister(screenshotHotkey);
+	obs_hotkey_unregister(sourceScreenshotHotkey);
 }
 
 OBSBasic::~OBSBasic()
@@ -2370,6 +2397,7 @@ OBSBasic::~OBSBasic()
 	if (updateCheckThread && updateCheckThread->isRunning())
 		updateCheckThread->wait();
 
+	delete screenshotData;
 	delete multiviewProjectorMenu;
 	delete previewProjector;
 	delete studioProgramProjector;
@@ -4330,6 +4358,8 @@ void OBSBasic::on_scenes_customContextMenuRequested(const QPoint &pos)
 			QTStr("SceneWindow"), this, SLOT(OpenSceneWindow()));
 
 		popup.addAction(sceneWindow);
+		popup.addAction(QTStr("Screenshot.Scene"), this,
+				SLOT(ScreenshotScene()));
 		popup.addSeparator();
 		popup.addAction(QTStr("Filters"), this,
 				SLOT(OpenSceneFilters()));
@@ -4681,6 +4711,9 @@ void OBSBasic::CreateSourcePopupMenu(int idx, bool preview)
 
 		popup.addAction(previewWindow);
 
+		popup.addAction(QTStr("Screenshot.Preview"), this,
+				SLOT(ScreenshotScene()));
+
 		popup.addSeparator();
 	}
 
@@ -4796,6 +4829,8 @@ void OBSBasic::CreateSourcePopupMenu(int idx, bool preview)
 
 		popup.addMenu(sourceProjector);
 		popup.addAction(sourceWindow);
+		popup.addAction(QTStr("Screenshot.Source"), this,
+				SLOT(ScreenshotSelectedSource()));
 		popup.addSeparator();
 
 		action = popup.addAction(QTStr("Interact"), this,
@@ -6249,6 +6284,9 @@ void OBSBasic::on_program_customContextMenuRequested(const QPoint &)
 
 	popup.addAction(studioProgramWindow);
 
+	popup.addAction(QTStr("Screenshot.StudioProgram"), this,
+			SLOT(ScreenshotProgram()));
+
 	popup.exec(QCursor::pos());
 }
 

UI/window-basic-main.hpp

@@ -432,6 +432,8 @@ private:
 	obs_hotkey_id togglePreviewProgramHotkey = 0;
 	obs_hotkey_id transitionHotkey = 0;
 	obs_hotkey_id statsHotkey = 0;
+	obs_hotkey_id screenshotHotkey = 0;
+	obs_hotkey_id sourceScreenshotHotkey = 0;
 	int quickTransitionIdCounter = 1;
 	bool overridingTransition = false;
 
@@ -525,6 +527,8 @@ private:
 	void UpdateProjectorHideCursor();
 	void UpdateProjectorAlwaysOnTop(bool top);
 
+	QPointer<QObject> screenshotData;
+
 public slots:
 	void DeferSaveBegin();
 	void DeferSaveEnd();
@@ -893,6 +897,10 @@ private slots:
 	void on_recordButton_clicked();
 	void VCamButtonClicked();
 	void on_settingsButton_clicked();
+	void Screenshot(OBSSource source_ = nullptr);
+	void ScreenshotSelectedSource();
+	void ScreenshotProgram();
+	void ScreenshotScene();
 
 	void on_actionHelpPortal_triggered();
 	void on_actionWebsite_triggered();

docs/sphinx/reference-frontend-api.rst

@@ -529,3 +529,17 @@ Functions
    active scene if not in studio mode.
 
    :param scene: The scene to set as the current preview.
+
+---------------------------------------
+
+.. function:: void *obs_frontend_take_screenshot(void)
+
+   Takes a screenshot of the main OBS output.
+
+---------------------------------------
+
+.. function:: void *obs_frontend_take_source_screenshot(obs_source_t *source)
+
+   Takes a screenshot of the specified source.
+
+   :param source: The source to take screenshot of.