NV12ToRgbShader

0_using.cpp

NV12ToRgbShader* shader = new NV12ToRgbShader();
if(shader->Init())
{
  AVFrame* src = ...;
  AVFrame* dst = ...;
  if(shader->Convert(src,dst))
  {
    
  }
}

NV12ToRgbShader.cpp

#include "pch.h"
#include "PixelShader.h"
#include "VertexShader.h"
#include "Utils.h"
#include "NV12ToRgbShader.h"
#include <math.h>

#define NUMVERTICES 6

typedef struct _VERTEX
{
	DirectX::XMFLOAT3 Pos;
	DirectX::XMFLOAT2 TexCoord;
} VERTEX;

// Vertices for drawing whole texture
VERTEX Vertices[NUMVERTICES] =
{
	{XMFLOAT3(-1.0f, -1.0f, 0), XMFLOAT2(0.0f, 1.0f)},
	{XMFLOAT3(-1.0f, 1.0f, 0), XMFLOAT2(0.0f, 0.0f)},
	{XMFLOAT3(1.0f, -1.0f, 0), XMFLOAT2(1.0f, 1.0f)},
	{XMFLOAT3(1.0f, -1.0f, 0), XMFLOAT2(1.0f, 1.0f)},
	{XMFLOAT3(-1.0f, 1.0f, 0), XMFLOAT2(0.0f, 0.0f)},
	{XMFLOAT3(1.0f, 1.0f, 0), XMFLOAT2(1.0f, 0.0f)},
};

FLOAT blendFactor[4] = { 0.f, 0.f, 0.f, 0.f };

NV12ToRgbShader::NV12ToRgbShader() {
}

NV12ToRgbShader::~NV12ToRgbShader() {
	this->ReleaseSharedSurf();
	_d3d11_vertexBuffer.Reset();
	_d3d11_pixelShader.Reset();
	_d3d11_inputLayout.Reset();
	_d3d11_vertexShader.Reset();
	_d3d11_samplerState.Reset();
	_d3d11_device.Reset();
	_d3d11_deviceCtx.Reset();
}

bool NV12ToRgbShader::Init() {
	HRESULT hr;

	// Driver types supported
	D3D_DRIVER_TYPE DriverTypes[] =
	{
		D3D_DRIVER_TYPE_HARDWARE,
		//D3D_DRIVER_TYPE_WARP,
		//D3D_DRIVER_TYPE_REFERENCE,
	};
	UINT NumDriverTypes = ARRAYSIZE(DriverTypes);

	// Feature levels supported
	D3D_FEATURE_LEVEL FeatureLevels[] =
	{
		D3D_FEATURE_LEVEL_11_0,
	};
	UINT NumFeatureLevels = ARRAYSIZE(FeatureLevels);
	D3D_FEATURE_LEVEL FeatureLevel;
	// This flag adds support for surfaces with a different color channel ordering
	// than the default. It is required for compatibility with Direct2D.
	UINT creationFlags = 
		D3D11_CREATE_DEVICE_SINGLETHREADED | 
		D3D11_CREATE_DEVICE_BGRA_SUPPORT;

	for (UINT DriverTypeIndex = 0; DriverTypeIndex < NumDriverTypes; ++DriverTypeIndex)
	{
		hr = D3D11CreateDevice(nullptr, DriverTypes[DriverTypeIndex], nullptr, creationFlags, FeatureLevels, NumFeatureLevels,
			D3D11_SDK_VERSION, this->_d3d11_device.GetAddressOf(), &FeatureLevel, this->_d3d11_deviceCtx.GetAddressOf());
		if (SUCCEEDED(hr))
		{
			// Device creation succeeded, no need to loop anymore
			break;
		}
	}
	if (FAILED(hr))
		return false;

	//SamplerState
	D3D11_SAMPLER_DESC desc = CD3D11_SAMPLER_DESC(CD3D11_DEFAULT());
	hr = this->_d3d11_device->CreateSamplerState(&desc, this->_d3d11_samplerState.GetAddressOf());
	if (FAILED(hr))
		return false;

	//VertexShader
	UINT Size = ARRAYSIZE(g_VS);
	hr = this->_d3d11_device->CreateVertexShader(g_VS, Size, nullptr, this->_d3d11_vertexShader.GetAddressOf());
	if (FAILED(hr))
		return false;
	constexpr std::array<D3D11_INPUT_ELEMENT_DESC, 2> Layout =
	{ {
		{ "POSITION", 0, DXGI_FORMAT_R32G32B32_FLOAT, 0, 0, D3D11_INPUT_PER_VERTEX_DATA, 0 },
		{ "TEXCOORD", 0, DXGI_FORMAT_R32G32_FLOAT,    0, 12, D3D11_INPUT_PER_VERTEX_DATA, 0 },
	} };

	hr = this->_d3d11_device->CreateInputLayout(Layout.data(), Layout.size(), g_VS, Size, this->_d3d11_inputLayout.GetAddressOf());
	if (FAILED(hr))
		return false;


	//PixelShader
	Size = ARRAYSIZE(g_PS);
	hr = this->_d3d11_device->CreatePixelShader(g_PS, Size, nullptr, this->_d3d11_pixelShader.GetAddressOf());
	if (FAILED(hr))
		return false;

	//VertexBuffer
	D3D11_BUFFER_DESC BufferDesc;
	RtlZeroMemory(&BufferDesc, sizeof(BufferDesc));
	BufferDesc.Usage = D3D11_USAGE_DEFAULT;
	BufferDesc.ByteWidth = sizeof(VERTEX) * NUMVERTICES;
	BufferDesc.BindFlags = D3D11_BIND_VERTEX_BUFFER;
	BufferDesc.CPUAccessFlags = 0;
	D3D11_SUBRESOURCE_DATA InitData;
	RtlZeroMemory(&InitData, sizeof(InitData));
	InitData.pSysMem = Vertices;
	hr = this->_d3d11_device->CreateBuffer(&BufferDesc, &InitData, this->_d3d11_vertexBuffer.GetAddressOf());
	if (FAILED(hr))
		return false;

	return true;
}


void NV12ToRgbShader::DeviceCtxSet(int width, int height) {
	//init set
	this->_d3d11_deviceCtx->IASetInputLayout(this->_d3d11_inputLayout.Get());
	this->_d3d11_deviceCtx->OMSetBlendState(nullptr, blendFactor, 0xffffffff);
	//this->_d3d11_deviceCtx->ClearRenderTargetView(this->_renderTargetView.Get(), blendFactor);
	this->_d3d11_deviceCtx->VSSetShader(this->_d3d11_vertexShader.Get(), nullptr, 0);
	this->_d3d11_deviceCtx->PSSetShader(this->_d3d11_pixelShader.Get(), nullptr, 0);
	this->_d3d11_deviceCtx->PSSetSamplers(0, 1, this->_d3d11_samplerState.GetAddressOf());
	this->_d3d11_deviceCtx->IASetPrimitiveTopology(D3D11_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
	UINT Stride = sizeof(VERTEX);
	UINT Offset = 0;
	this->_d3d11_deviceCtx->IASetVertexBuffers(0, 1, this->_d3d11_vertexBuffer.GetAddressOf(), &Stride, &Offset);


	//SharedSurf	
	std::array<ID3D11ShaderResourceView*, 2> const textureViews = {
		this->_luminanceView.Get(),
		this->_chrominanceView.Get()
	};
	this->_d3d11_deviceCtx->PSSetShaderResources(0, textureViews.size(), textureViews.data());
	this->_d3d11_deviceCtx->OMSetRenderTargets(1, this->_renderTargetView.GetAddressOf(), nullptr);

	D3D11_VIEWPORT VP;
	VP.Width = static_cast<FLOAT>(width);
	VP.Height = static_cast<FLOAT>(height);
	VP.MinDepth = 0.0f;
	VP.MaxDepth = 1.0f;
	VP.TopLeftX = 0;
	VP.TopLeftY = 0;
	this->_d3d11_deviceCtx->RSSetViewports(1, &VP);
	//this->_d3d11_deviceCtx->Dispatch(8, 8, 1);
	this->_d3d11_deviceCtx->Dispatch(
		(UINT)ceil(width * 1.0 / 8),
		(UINT)ceil(height * 1.0 / 8),
		1);
	this->_width = width;
	this->_height = height;
}

bool NV12ToRgbShader::CreateSharedSurf(int width, int height) {
	//
	HRESULT hr{ 0 };

	D3D11_TEXTURE2D_DESC texDesc_nv12;
	ZeroMemory(&texDesc_nv12, sizeof(texDesc_nv12));
	texDesc_nv12.Format = DXGI_FORMAT_NV12;
	texDesc_nv12.Width = width;
	texDesc_nv12.Height = height;
	texDesc_nv12.ArraySize = 1;
	texDesc_nv12.MipLevels = 1;
	texDesc_nv12.BindFlags = D3D11_BIND_SHADER_RESOURCE;
	texDesc_nv12.Usage = D3D11_USAGE_DYNAMIC;
	texDesc_nv12.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
	texDesc_nv12.SampleDesc.Count = 1;
	texDesc_nv12.SampleDesc.Quality = 0;
	texDesc_nv12.MiscFlags = 0;
	hr = this->_d3d11_device->CreateTexture2D(&texDesc_nv12, nullptr, this->_texture_nv12.GetAddressOf());
	if (FAILED(hr))
		return false;

	//
	D3D11_SHADER_RESOURCE_VIEW_DESC const luminancePlaneDesc
		= CD3D11_SHADER_RESOURCE_VIEW_DESC(this->_texture_nv12.Get(), D3D11_SRV_DIMENSION_TEXTURE2D, DXGI_FORMAT_R8_UNORM);
	hr = this->_d3d11_device->CreateShaderResourceView(this->_texture_nv12.Get(), &luminancePlaneDesc, this->_luminanceView.GetAddressOf());
	if (FAILED(hr))
		return false;

	//
	D3D11_SHADER_RESOURCE_VIEW_DESC const chrominancePlaneDesc
		= CD3D11_SHADER_RESOURCE_VIEW_DESC(this->_texture_nv12.Get(), D3D11_SRV_DIMENSION_TEXTURE2D, DXGI_FORMAT_R8G8_UNORM);
	hr = this->_d3d11_device->CreateShaderResourceView(this->_texture_nv12.Get(), &chrominancePlaneDesc, this->_chrominanceView.GetAddressOf());
	if (FAILED(hr))
		return false;


	D3D11_TEXTURE2D_DESC texDesc_rgba;
	ZeroMemory(&texDesc_rgba, sizeof(texDesc_rgba));
	texDesc_rgba.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
	texDesc_rgba.Width = width;
	texDesc_rgba.Height = height;
	texDesc_rgba.ArraySize = 1;
	texDesc_rgba.MipLevels = 1;
	texDesc_rgba.BindFlags = D3D11_BIND_RENDER_TARGET | D3D11_BIND_SHADER_RESOURCE;
	texDesc_rgba.Usage = D3D11_USAGE_DEFAULT;
	texDesc_rgba.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
	texDesc_rgba.SampleDesc.Count = 1;
	texDesc_rgba.SampleDesc.Quality = 0;
	texDesc_rgba.MiscFlags = 0;
	hr = this->_d3d11_device->CreateTexture2D(&texDesc_rgba, nullptr, this->_texture_rgba_target.GetAddressOf());
	if (FAILED(hr))
		return false;

	texDesc_rgba.BindFlags = 0;
	texDesc_rgba.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
	texDesc_rgba.Usage = D3D11_USAGE_STAGING;//cpu read
	texDesc_rgba.MiscFlags = 0;
	hr = this->_d3d11_device->CreateTexture2D(&texDesc_rgba, nullptr, this->_texture_rgba_copy.GetAddressOf());
	if (FAILED(hr))
		return false;


	D3D11_RENDER_TARGET_VIEW_DESC rtvDesc{};
	rtvDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
	rtvDesc.ViewDimension = D3D11_RTV_DIMENSION_TEXTURE2D;
	rtvDesc.Texture2D.MipSlice = 0;
	hr = this->_d3d11_device->CreateRenderTargetView(this->_texture_rgba_target.Get(), &rtvDesc, this->_renderTargetView.GetAddressOf());
	if (FAILED(hr))
		return false;

	return true;
}

void NV12ToRgbShader::ReleaseSharedSurf() {
	this->_d3d11_deviceCtx->ClearState();
	this->_texture_nv12.Reset();
	this->_luminanceView.Reset();
	this->_chrominanceView.Reset();
	this->_texture_rgba_target.Reset();
	this->_texture_rgba_copy.Reset();
	this->_renderTargetView.Reset();
	this->_width = 0;
	this->_height = 0;
}

//https://medium.com/swlh/streaming-video-with-ffmpeg-and-directx-11-7395fcb372c4
bool NV12ToRgbShader::Convert(const AVFrame* source, AVFrame* received) {
	HRESULT hr{ 0 };
	if (source == NULL || received == NULL)
		return false;
	if (source->format != AV_PIX_FMT_D3D11)
		return false;
	if (!source->hw_frames_ctx)
		return false;

	//init/reinit shader surface
	if (this->_width != source->width || this->_height != source->height) {
		this->ReleaseSharedSurf();
		if (this->CreateSharedSurf(source->width, source->height))
		    this->DeviceCtxSet(source->width, source->height);
		else
		    return false;
	}

	ComPtr<ID3D11Texture2D> texture = (ID3D11Texture2D*)source->data[0];
	const int texture_index = (int)source->data[1];
	//bind/copy ffmpeg hw texture -> local d3d11 texture
	this->_d3d11_deviceCtx->CopySubresourceRegion(
		this->_texture_nv12.Get(), 0, 0, 0, 0,
		texture.Get(), texture_index, nullptr
	);

	this->_d3d11_deviceCtx->Draw(NUMVERTICES, 0);


	//render target view only 1 sub resource https://docs.microsoft.com/en-us/windows/win32/direct3d11/overviews-direct3d-11-resources-subresources
	#define CopySubResource 0
	this->_d3d11_deviceCtx->CopyResource(this->_texture_rgba_copy.Get(), this->_texture_rgba_target.Get());	

	//get texture output
	D3D11_MAPPED_SUBRESOURCE ms;
	hr = this->_d3d11_deviceCtx->Map(this->_texture_rgba_copy.Get(), CopySubResource, D3D11_MAP_READ, 0, &ms);
	if (FAILED(hr))
		return false;

	bool result = this->CopyMapResource(ms, source, received);

	this->_d3d11_deviceCtx->Unmap(this->_texture_rgba_copy.Get(), 0);

	return result;
}

bool NV12ToRgbShader::CopyMapResource(const D3D11_MAPPED_SUBRESOURCE& ms, const AVFrame* source, AVFrame* received) {
	bool result = false;
	int size = av_image_get_buffer_size(AVPixelFormat::AV_PIX_FMT_BGRA, source->width, source->height, 1);
	if (size <= ms.DepthPitch)
	{
		av_frame_unref(received);
		AVBufferRef* dataref = av_buffer_alloc(size);
		if (dataref != nullptr &&
			avcheck(av_image_fill_arrays(
				received->data, received->linesize, dataref->data,
				AVPixelFormat::AV_PIX_FMT_BGRA, source->width, source->height, 1)))
		{
			av_frame_copy_props(received, source);
			received->format = AVPixelFormat::AV_PIX_FMT_BGRA;
			received->width = source->width;
			received->height = source->height;
			received->pts = source->pts;
			received->pkt_dts = source->pkt_dts;
			received->time_base = source->time_base;
			received->pkt_duration = source->pkt_duration;
			received->pkt_pos = source->pkt_pos;

			received->buf[0] = dataref;
			if (received->linesize[0] == ms.RowPitch)
			{
				memcpy(dataref->data, ms.pData, ms.DepthPitch);
			}
			else
			{
				for (UINT64 i = 0; i < source->height; i++)
				{
					uint8_t* dst = dataref->data + i * received->linesize[0];
					uint8_t* src = (uint8_t*)ms.pData + i * ms.RowPitch;
					memcpy(dst, src, received->linesize[0]);
				}
			}
			result = true;
		}
		else
		{
			av_buffer_unref(&dataref);
		}
	}
	return result;
}

NV12ToRgbShader.h

#ifndef NV12ToRgbShader_H
#define NV12ToRgbShader_H
class NV12ToRgbShader
{
public:
	NV12ToRgbShader();
	~NV12ToRgbShader();

	bool Init();

	bool Convert(const AVFrame* source, AVFrame* received);
private:
	//init
	ComPtr<ID3D11DeviceContext> _d3d11_deviceCtx = nullptr;
	ComPtr<ID3D11Device> _d3d11_device = nullptr;
	ComPtr<ID3D11SamplerState> _d3d11_samplerState = nullptr;
	ComPtr<ID3D11VertexShader> _d3d11_vertexShader = nullptr;
	ComPtr<ID3D11InputLayout> _d3d11_inputLayout = nullptr;
	ComPtr<ID3D11PixelShader> _d3d11_pixelShader = nullptr;
	ComPtr<ID3D11Buffer> _d3d11_vertexBuffer = nullptr;


	void DeviceCtxSet(int width, int height);

	

	//SharedSurf
	
	ComPtr<ID3D11Texture2D> _texture_nv12 = nullptr;
	ComPtr<ID3D11ShaderResourceView> _luminanceView = nullptr;
	ComPtr<ID3D11ShaderResourceView> _chrominanceView = nullptr;
	ComPtr<ID3D11RenderTargetView> _renderTargetView = nullptr;
	ComPtr<ID3D11Texture2D> _texture_rgba_target = nullptr;
	ComPtr<ID3D11Texture2D> _texture_rgba_copy = nullptr;
	uint32_t _width{ 0 };
	uint32_t _height{ 0 };
	bool CreateSharedSurf(int width, int height);
	void ReleaseSharedSurf();
	bool CopyMapResource(const D3D11_MAPPED_SUBRESOURCE& ms, const AVFrame* source, AVFrame* received);
};
#endif // !NV12ToRgbShader_H

PixelShader.hlsl

//*********************************************************
//
// Copyright (c) Microsoft. All rights reserved.
// This code is licensed under the MIT License (MIT).
// THIS CODE IS PROVIDED *AS IS* WITHOUT WARRANTY OF
// ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING ANY
// IMPLIED WARRANTIES OF FITNESS FOR A PARTICULAR
// PURPOSE, MERCHANTABILITY, OR NON-INFRINGEMENT.
//
//*********************************************************

// Per-pixel color data passed through the pixel shader.
struct PixelShaderInput
{
	float4 pos         : SV_POSITION;
	float2 texCoord    : TEXCOORD0;
};

Texture2D<float>  luminanceChannel   : t0;
Texture2D<float2> chrominanceChannel : t1;
SamplerState      defaultSampler     : s0;

// Derived from https://msdn.microsoft.com/en-us/library/windows/desktop/dd206750(v=vs.85).aspx
// Section: Converting 8-bit YUV to RGB888
static const float3x3 YUVtoRGBCoeffMatrix =
{
	1.164383f,  1.164383f, 1.164383f,
	0.000000f, -0.391762f, 2.017232f,
	1.596027f, -0.812968f, 0.000000f
};

float3 ConvertYUVtoRGB(float3 yuv)
{
	// Derived from https://msdn.microsoft.com/en-us/library/windows/desktop/dd206750(v=vs.85).aspx
	// Section: Converting 8-bit YUV to RGB888

	// These values are calculated from (16 / 255) and (128 / 255)
	yuv -= float3(0.062745f, 0.501960f, 0.501960f);
	yuv = mul(yuv, YUVtoRGBCoeffMatrix);
	yuv = saturate(yuv);//BGR
	return float3(yuv.z, yuv.y, yuv.x);
}

[numthreads (8,8,1)]
float4 PS(PixelShaderInput input) : SV_TARGET
{
	float y = luminanceChannel.Sample(defaultSampler, input.texCoord);
	float2 uv = chrominanceChannel.Sample(defaultSampler, input.texCoord);

	return float4(ConvertYUVtoRGB(float3(y, uv)), 1.f);
}

VertexShader.hlsl

// THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF
// ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A
// PARTICULAR PURPOSE.
//
// Copyright (c) Microsoft Corporation. All rights reserved
//----------------------------------------------------------------------

struct VS_INPUT
{
    float4 Pos : POSITION;
    float2 Tex : TEXCOORD;
};

struct VS_OUTPUT
{
    float4 Pos : SV_POSITION;
    float2 Tex : TEXCOORD;
};


//--------------------------------------------------------------------------------------
// Vertex Shader
//--------------------------------------------------------------------------------------
VS_OUTPUT VS(VS_INPUT input)
{
    return input;
}