138 lines
3.7 KiB
C++
138 lines
3.7 KiB
C++
#include "PngScreenshotWriter.h"
|
|
|
|
#include <windows.h>
|
|
#include <wincodec.h>
|
|
#include <atlbase.h>
|
|
|
|
#include <sstream>
|
|
#include <thread>
|
|
|
|
namespace
|
|
{
|
|
std::string HResultToString(HRESULT hr)
|
|
{
|
|
std::ostringstream stream;
|
|
stream << "HRESULT 0x" << std::hex << static_cast<unsigned long>(hr);
|
|
return stream.str();
|
|
}
|
|
|
|
bool WritePngFile(
|
|
const std::filesystem::path& outputPath,
|
|
unsigned width,
|
|
unsigned height,
|
|
const std::vector<unsigned char>& bgraPixels,
|
|
std::string& error)
|
|
{
|
|
if (width == 0 || height == 0 || bgraPixels.size() < static_cast<std::size_t>(width) * height * 4)
|
|
{
|
|
error = "Invalid screenshot dimensions or pixel buffer.";
|
|
return false;
|
|
}
|
|
|
|
HRESULT initializeResult = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
|
|
const bool shouldUninitialize = SUCCEEDED(initializeResult);
|
|
if (FAILED(initializeResult) && initializeResult != RPC_E_CHANGED_MODE)
|
|
{
|
|
error = "CoInitializeEx failed: " + HResultToString(initializeResult);
|
|
return false;
|
|
}
|
|
|
|
CComPtr<IWICImagingFactory> factory;
|
|
HRESULT result = CoCreateInstance(
|
|
CLSID_WICImagingFactory,
|
|
nullptr,
|
|
CLSCTX_INPROC_SERVER,
|
|
IID_PPV_ARGS(&factory));
|
|
if (FAILED(result))
|
|
{
|
|
error = "Could not create WIC imaging factory: " + HResultToString(result);
|
|
if (shouldUninitialize)
|
|
CoUninitialize();
|
|
return false;
|
|
}
|
|
|
|
CComPtr<IWICStream> stream;
|
|
result = factory->CreateStream(&stream);
|
|
if (SUCCEEDED(result))
|
|
result = stream->InitializeFromFilename(outputPath.wstring().c_str(), GENERIC_WRITE);
|
|
if (FAILED(result))
|
|
{
|
|
error = "Could not open screenshot output file: " + HResultToString(result);
|
|
if (shouldUninitialize)
|
|
CoUninitialize();
|
|
return false;
|
|
}
|
|
|
|
CComPtr<IWICBitmapEncoder> encoder;
|
|
result = factory->CreateEncoder(GUID_ContainerFormatPng, nullptr, &encoder);
|
|
if (SUCCEEDED(result))
|
|
result = encoder->Initialize(stream, WICBitmapEncoderNoCache);
|
|
if (FAILED(result))
|
|
{
|
|
error = "Could not initialize PNG encoder: " + HResultToString(result);
|
|
if (shouldUninitialize)
|
|
CoUninitialize();
|
|
return false;
|
|
}
|
|
|
|
CComPtr<IWICBitmapFrameEncode> frame;
|
|
CComPtr<IPropertyBag2> propertyBag;
|
|
result = encoder->CreateNewFrame(&frame, &propertyBag);
|
|
if (SUCCEEDED(result))
|
|
result = frame->Initialize(propertyBag);
|
|
if (SUCCEEDED(result))
|
|
result = frame->SetSize(width, height);
|
|
|
|
WICPixelFormatGUID pixelFormat = GUID_WICPixelFormat32bppBGRA;
|
|
if (SUCCEEDED(result))
|
|
result = frame->SetPixelFormat(&pixelFormat);
|
|
if (SUCCEEDED(result) && pixelFormat != GUID_WICPixelFormat32bppBGRA)
|
|
{
|
|
error = "PNG encoder did not accept BGRA pixel format.";
|
|
result = E_FAIL;
|
|
}
|
|
|
|
const UINT stride = width * 4;
|
|
const UINT imageSize = stride * height;
|
|
if (SUCCEEDED(result))
|
|
result = frame->WritePixels(height, stride, imageSize, const_cast<BYTE*>(bgraPixels.data()));
|
|
if (SUCCEEDED(result))
|
|
result = frame->Commit();
|
|
if (SUCCEEDED(result))
|
|
result = encoder->Commit();
|
|
|
|
if (shouldUninitialize)
|
|
CoUninitialize();
|
|
|
|
if (FAILED(result))
|
|
{
|
|
error = "Could not write screenshot PNG: " + HResultToString(result);
|
|
std::error_code ignored;
|
|
std::filesystem::remove(outputPath, ignored);
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
}
|
|
|
|
void WritePngFileAsync(
|
|
const std::filesystem::path& outputPath,
|
|
unsigned width,
|
|
unsigned height,
|
|
std::vector<unsigned char> rgbaPixels)
|
|
{
|
|
std::thread(
|
|
[outputPath, width, height, pixels = std::move(rgbaPixels)]() mutable
|
|
{
|
|
for (std::size_t index = 0; index + 3 < pixels.size(); index += 4)
|
|
std::swap(pixels[index], pixels[index + 2]);
|
|
|
|
std::string error;
|
|
if (!WritePngFile(outputPath, width, height, pixels, error))
|
|
OutputDebugStringA(("Screenshot write failed: " + error + "\n").c_str());
|
|
else
|
|
OutputDebugStringA(("Screenshot written: " + outputPath.string() + "\n").c_str());
|
|
}).detach();
|
|
}
|