Skip to content

Instantly share code, notes, and snippets.

@tqk2811
Last active May 13, 2022
Embed
What would you like to do?
NV12ToRgbShader
NV12ToRgbShader* shader = new NV12ToRgbShader();
if(shader->Init())
{
AVFrame* src = ...;
AVFrame* dst = ...;
if(shader->Convert(src,dst))
{
}
}
#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;
}
#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
//*********************************************************
//
// 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);
}
// 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;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment