UI: JXR screenshots on Windows

Use JXR for HDR video on Windows. Other operating systems will tonemap
HDR to SDR, and save to PNG. Continue to take PNG screenshots for SDR.
We will probably support EXR for Mac/Linux someday.

UI/screenshot-obj.hpp

@@ -36,6 +36,7 @@ public:
 	OBSWeakSource weakSource;
 	std::string path;
 	QImage image;
+	std::vector<uint8_t> half_bytes;
 	uint32_t cx;
 	uint32_t cy;
 	std::thread th;

UI/window-basic-main-screenshot.cpp

@@ -19,6 +19,13 @@
 #include "screenshot-obj.hpp"
 #include "qt-wrappers.hpp"
 
+#ifdef _WIN32
+#include <wincodec.h>
+#include <wincodecsdk.h>
+#include <wrl/client.h>
+#pragma comment(lib, "windowscodecs.lib")
+#endif
+
 static void ScreenshotTick(void *param, float);
 
 /* ========================================================================= */
@@ -71,11 +78,23 @@ void ScreenshotObj::Screenshot()
 		return;
 	}
 
-	texrender = gs_texrender_create(GS_RGBA, GS_ZS_NONE);
-	stagesurf = gs_stagesurface_create(cx, cy, GS_RGBA);
+#ifdef _WIN32
+	enum gs_color_space space =
+		obs_source_get_color_space(source, 0, nullptr);
+	if (space == GS_CS_709_EXTENDED) {
+		/* Convert for JXR */
+		space = GS_CS_709_SCRGB;
+	}
+#else
+	/* Tonemap to SDR if HDR */
+	const enum gs_color_space space = GS_CS_SRGB;
+#endif
+	const enum gs_color_format format = gs_get_format_from_space(space);
 
-	gs_texrender_reset(texrender);
-	if (gs_texrender_begin(texrender, cx, cy)) {
+	texrender = gs_texrender_create(format, GS_ZS_NONE);
+	stagesurf = gs_stagesurface_create(cx, cy, format);
+
+	if (gs_texrender_begin_with_color_space(texrender, cx, cy, space)) {
 		vec4 zero;
 		vec4_zero(&zero);
 
@@ -108,13 +127,26 @@ void ScreenshotObj::Copy()
 	uint8_t *videoData = nullptr;
 	uint32_t videoLinesize = 0;
 
-	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);
+		if (gs_stagesurface_get_color_format(stagesurf) == GS_RGBA16F) {
+			const uint32_t linesize = cx * 8;
+			half_bytes.reserve(cx * cy * 8);
+
+			for (uint32_t y = 0; y < cy; y++) {
+				const uint8_t *const line =
+					videoData + (y * videoLinesize);
+				half_bytes.insert(half_bytes.end(), line,
+						  line + linesize);
+			}
+		} else {
+			image = QImage(cx, cy, QImage::Format::Format_RGBX8888);
+
+			int linesize = image.bytesPerLine();
+			for (int y = 0; y < (int)cy; y++)
+				memcpy(image.scanLine(y),
+				       videoData + (y * videoLinesize),
+				       linesize);
+		}
 
 		gs_stagesurface_unmap(stagesurf);
 	}
@@ -143,17 +175,122 @@ void ScreenshotObj::Save()
 	bool overwriteIfExists =
 		config_get_bool(config, "Output", "OverwriteIfExists");
 
+	const char *ext = half_bytes.empty() ? "png" : "jxr";
 	path = GetOutputFilename(
-		rec_path, "png", noSpace, overwriteIfExists,
+		rec_path, ext, noSpace, overwriteIfExists,
 		GetFormatString(filenameFormat, "Screenshot", nullptr).c_str());
 
 	th = std::thread([this] { MuxAndFinish(); });
 }
 
+#ifdef _WIN32
+static HRESULT SaveJxrImage(LPCWSTR path, uint8_t *pixels, uint32_t cx,
+			    uint32_t cy, IWICBitmapFrameEncode *frameEncode,
+			    IPropertyBag2 *options)
+{
+	wchar_t lossless[] = L"Lossless";
+	PROPBAG2 bag = {};
+	bag.pstrName = lossless;
+	VARIANT value = {};
+	value.vt = VT_BOOL;
+	value.bVal = TRUE;
+	HRESULT hr = options->Write(1, &bag, &value);
+	if (FAILED(hr))
+		return hr;
+
+	hr = frameEncode->Initialize(options);
+	if (FAILED(hr))
+		return hr;
+
+	hr = frameEncode->SetSize(cx, cy);
+	if (FAILED(hr))
+		return hr;
+
+	hr = frameEncode->SetResolution(72, 72);
+	if (FAILED(hr))
+		return hr;
+
+	WICPixelFormatGUID pixelFormat = GUID_WICPixelFormat64bppRGBAHalf;
+	hr = frameEncode->SetPixelFormat(&pixelFormat);
+	if (FAILED(hr))
+		return hr;
+
+	if (memcmp(&pixelFormat, &GUID_WICPixelFormat64bppRGBAHalf,
+		   sizeof(WICPixelFormatGUID)) != 0)
+		return E_FAIL;
+
+	hr = frameEncode->WritePixels(cy, cx * 8, cx * cy * 8, pixels);
+	if (FAILED(hr))
+		return hr;
+
+	hr = frameEncode->Commit();
+	if (FAILED(hr))
+		return hr;
+
+	return S_OK;
+}
+
+static HRESULT SaveJxr(LPCWSTR path, uint8_t *pixels, uint32_t cx, uint32_t cy)
+{
+	Microsoft::WRL::ComPtr<IWICImagingFactory> factory;
+	HRESULT hr = CoCreateInstance(CLSID_WICImagingFactory, NULL,
+				      CLSCTX_INPROC_SERVER,
+				      IID_PPV_ARGS(factory.GetAddressOf()));
+	if (FAILED(hr))
+		return hr;
+
+	Microsoft::WRL::ComPtr<IWICStream> stream;
+	hr = factory->CreateStream(stream.GetAddressOf());
+	if (FAILED(hr))
+		return hr;
+
+	hr = stream->InitializeFromFilename(path, GENERIC_WRITE);
+	if (FAILED(hr))
+		return hr;
+
+	Microsoft::WRL::ComPtr<IWICBitmapEncoder> encoder = NULL;
+	hr = factory->CreateEncoder(GUID_ContainerFormatWmp, NULL,
+				    encoder.GetAddressOf());
+	if (FAILED(hr))
+		return hr;
+
+	hr = encoder->Initialize(stream.Get(), WICBitmapEncoderNoCache);
+	if (FAILED(hr))
+		return hr;
+
+	Microsoft::WRL::ComPtr<IWICBitmapFrameEncode> frameEncode;
+	Microsoft::WRL::ComPtr<IPropertyBag2> options;
+	hr = encoder->CreateNewFrame(frameEncode.GetAddressOf(),
+				     options.GetAddressOf());
+	if (FAILED(hr))
+		return hr;
+
+	hr = SaveJxrImage(path, pixels, cx, cy, frameEncode.Get(),
+			  options.Get());
+	if (FAILED(hr))
+		return hr;
+
+	encoder->Commit();
+	return S_OK;
+}
+#endif // #ifdef _WIN32
+
 void ScreenshotObj::MuxAndFinish()
 {
-	image.save(QT_UTF8(path.c_str()));
-	blog(LOG_INFO, "Saved screenshot to '%s'", path.c_str());
+	if (half_bytes.empty()) {
+		image.save(QT_UTF8(path.c_str()));
+		blog(LOG_INFO, "Saved screenshot to '%s'", path.c_str());
+	} else {
+#ifdef _WIN32
+		wchar_t *path_w = nullptr;
+		os_utf8_to_wcs_ptr(path.c_str(), 0, &path_w);
+		if (path_w) {
+			SaveJxr(path_w, half_bytes.data(), cx, cy);
+			bfree(path_w);
+		}
+#endif // #ifdef _WIN32
+	}
+
 	deleteLater();
 }