[WIN32] DX11 DXGI Screen capture sample

dx11-screen-capture.cpp

// WIN32/C++17: DX11 DXGI Screen capture sample
//
// References:
//  - https://github.com/microsoftarchive/msdn-code-gallery-microsoft/tree/master/Official%20Windows%20Platform%20Sample/DXGI%20desktop%20duplication%20sample
//  - https://github.com/microsoft/DirectXTex/blob/master/ScreenGrab/ScreenGrab11.cpp
//  - https://docs.microsoft.com/en-us/windows/win32/direct3ddxgi/desktop-dup-api
//
#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <atlbase.h>
#include <dxgi1_2.h>
#include <d3d11.h>
#include <vector>

#pragma comment(lib, "D3D11.lib")

#define Report(x) { printf("Report: %s(%d), hr=0x%08x\n", __FILE__, __LINE__, (x)); }


struct Dev {
    CComPtr<ID3D11Device> device;
    CComPtr<ID3D11DeviceContext> deviceContext;
    D3D_FEATURE_LEVEL featureLevel;

    Dev() {
        static const D3D_DRIVER_TYPE driverTypes[] = {
            D3D_DRIVER_TYPE_HARDWARE,
            D3D_DRIVER_TYPE_WARP,
            D3D_DRIVER_TYPE_REFERENCE
        };

        static const D3D_FEATURE_LEVEL featureLevels[] = {
            D3D_FEATURE_LEVEL_11_0,
            D3D_FEATURE_LEVEL_10_1,
            D3D_FEATURE_LEVEL_10_0,
            D3D_FEATURE_LEVEL_9_1
        };

        for(const auto& driverType : driverTypes) {
            const auto hr = D3D11CreateDevice(
                nullptr,
                driverType,
                nullptr,
                0,
                featureLevels,
                static_cast<UINT>(std::size(featureLevels)),
                D3D11_SDK_VERSION,
                &device,
                &featureLevel,
                &deviceContext
            );
            if(SUCCEEDED(hr)) {
                break;
            }
            device.Release();
            deviceContext.Release();
        }
    }
};


struct OutputDuplication {
    CComPtr<IDXGIOutputDuplication> outputDuplication;

    OutputDuplication(ID3D11Device* device) {
        HRESULT hr;

        CComPtr<IDXGIDevice> dxgiDevice;
        hr = device->QueryInterface(
            __uuidof(dxgiDevice),
            reinterpret_cast<void**>(&dxgiDevice)
        );
        if(FAILED(hr)) { Report(hr); return; }

        CComPtr<IDXGIAdapter> dxgiAdapter;
        hr = dxgiDevice->GetParent(
            __uuidof(dxgiAdapter),
            reinterpret_cast<void**>(&dxgiAdapter)
        );
        if(FAILED(hr)) { Report(hr); return; }

        CComPtr<IDXGIOutput> dxgiOutput;
        hr = dxgiAdapter->EnumOutputs(0, &dxgiOutput);
        if(FAILED(hr)) { Report(hr); return; }

        CComPtr<IDXGIOutput1> dxgiOutput1;
        hr = dxgiOutput->QueryInterface(
            __uuidof(IDXGIOutput1),
            reinterpret_cast<void**>(&dxgiOutput1)
        );
        if(FAILED(hr)) { Report(hr); return; }

        hr = dxgiOutput1->DuplicateOutput(device, &outputDuplication);
        if(FAILED(hr)) { Report(hr); return; }
    }
};


struct AcquiredDesktopImage {
    CComPtr<ID3D11Texture2D> acquiredDesktopImage;

    AcquiredDesktopImage(IDXGIOutputDuplication* outputDuplication) {
        CComPtr<IDXGIResource> desktopResource;
        HRESULT hr = E_FAIL;
        for(int i = 0; i < 10; ++i) {
            DXGI_OUTDUPL_FRAME_INFO fi {};
            const int timeoutMsec = 500; // milliseconds
            hr = outputDuplication->AcquireNextFrame(timeoutMsec, &fi, &desktopResource);
            if(SUCCEEDED(hr) && (fi.LastPresentTime.QuadPart == 0)) {
                // If AcquireNextFrame() returns S_OK and
                // fi.LastPresentTime.QuadPart == 0, it means
                // AcquireNextFrame() didn't acquire next frame yet.
                // We must wait next frame sync timing to retrieve
                // actual frame data.
                //
                // Since method is successfully completed,
                // we need to release the resource and frame explicitly.
                desktopResource.Release();
                outputDuplication->ReleaseFrame();
                Sleep(1);
                continue;
            } else {
                break;
            }
        }
        if(FAILED(hr)) { Report(hr); return; }

        hr = desktopResource->QueryInterface(
            __uuidof(ID3D11Texture2D),
            reinterpret_cast<void**>(&acquiredDesktopImage)
        );
        if(FAILED(hr)) { Report(hr); return; }
    }
};


struct Image {
    std::vector<byte> bytes;
    int width = 0;
    int height = 0;
    int rowPitch = 0;
};


Image captureDesktop() {
    Dev dev;
    ID3D11Device* device = dev.device;
    ID3D11DeviceContext* deviceContext = dev.deviceContext;
    if(device == nullptr) { Report(E_FAIL); return {}; }

    // Create tex2dStaging which represents duplication image of desktop.
    CComPtr<ID3D11Texture2D> tex2dStaging;
    {
        OutputDuplication od(device);
        IDXGIOutputDuplication* outputDuplication = od.outputDuplication;
        if(outputDuplication == nullptr) { Report(E_FAIL); return {}; }

        AcquiredDesktopImage adi(outputDuplication);
        ID3D11Texture2D* acquiredDesktopImage = adi.acquiredDesktopImage;
        if(acquiredDesktopImage == nullptr) { Report(E_FAIL); return {}; }

        DXGI_OUTDUPL_DESC duplDesc;
        outputDuplication->GetDesc(&duplDesc);

        const auto f = static_cast<int>(duplDesc.ModeDesc.Format);
        const auto goodFormat = f == DXGI_FORMAT_B8G8R8A8_UNORM
                             || f == DXGI_FORMAT_B8G8R8X8_UNORM
                             || f == DXGI_FORMAT_B8G8R8A8_TYPELESS
                             || f == DXGI_FORMAT_B8G8R8A8_UNORM_SRGB
                             || f == DXGI_FORMAT_B8G8R8X8_TYPELESS
                             || f == DXGI_FORMAT_B8G8R8X8_UNORM_SRGB;
        if(! goodFormat) { Report(E_FAIL); return {}; }

        D3D11_TEXTURE2D_DESC desc {};
        desc.Width              = duplDesc.ModeDesc.Width;
        desc.Height             = duplDesc.ModeDesc.Height;
        desc.Format             = duplDesc.ModeDesc.Format;
        desc.ArraySize          = 1;
        desc.BindFlags          = 0;
        desc.MiscFlags          = 0;
        desc.SampleDesc.Count   = 1;
        desc.SampleDesc.Quality = 0;
        desc.MipLevels          = 1;
        desc.CPUAccessFlags     = D3D11_CPU_ACCESS_READ;
        desc.Usage              = D3D11_USAGE_STAGING;
        const auto hr = device->CreateTexture2D(&desc, nullptr, &tex2dStaging);
        if(FAILED(hr)) { Report(hr); return {}; }
        if(tex2dStaging == nullptr) { Report(E_FAIL); return {}; }

        deviceContext->CopyResource(tex2dStaging, acquiredDesktopImage);
    }

    // Lock tex2dStaging and copy its content from GPU to CPU memory.
    Image image;

    D3D11_TEXTURE2D_DESC desc;
    tex2dStaging->GetDesc(&desc);

    D3D11_MAPPED_SUBRESOURCE res;
    const auto hr = deviceContext->Map(
        tex2dStaging,
        D3D11CalcSubresource(0, 0, 0),
        D3D11_MAP_READ,
        0,
        &res
    );
    if(FAILED(hr)) { Report(hr); return {}; }
    image.width     = static_cast<int>(desc.Width);
    image.height    = static_cast<int>(desc.Height);
    image.rowPitch  = res.RowPitch;
    image.bytes.resize(image.rowPitch * image.height);
    memcpy(image.bytes.data(), res.pData, image.bytes.size());
    deviceContext->Unmap(tex2dStaging, 0);
    return image;
}


int main() {
    bool result = false;
    const auto image = captureDesktop();
    if(! image.bytes.empty()) {
        const char* filename = "screenshot.ppm";
        FILE* fp;
        if(fopen_s(&fp, filename, "wb") == 0) {
            // PPM format: https://en.wikipedia.org/wiki/Netpbm
            fprintf(fp, "P6\n#\n%d %d %d\n", image.width, image.height, 255);
            for(int y = 0; y < image.height; ++y) {
                for(int x = 0; x < image.width; ++x) {
                    const auto* p = image.bytes.data() + (image.rowPitch * y) + (4 * x);
                    fputc(p[2], fp);    //R
                    fputc(p[1], fp);    //G
                    fputc(p[0], fp);    //B
                }
            }
            fclose(fp);
            result = true;
        }
    }
    printf(result ? "OK\n" : "Fail\n");
}