Skip to content

Instantly share code, notes, and snippets.

@kebby
Last active April 26, 2020 06:03
Show Gist options
  • Save kebby/8cc8447619899b33e47fba836dd00ae2 to your computer and use it in GitHub Desktop.
Save kebby/8cc8447619899b33e47fba836dd00ae2 to your computer and use it in GitHub Desktop.
// Media Foundation h.264 encode test, by Tammo "kb" Hinrichs, 2016-8-31
// This source code file is in the public domain.
// This really is a minimal test, and it only does the raw encoding, no
// color space conversion or muxing whatsoever. So to make it usable you'll
// probably want to put the main loop into its own thread and make it poll
// the rendering thread for new image data instead of the RenderImage() call,
// and you'll probably also want to construct an MF graphfor color space
// conversion/muxing/audio, or call the respective MFTs manually. And get rid
// of the memory buffers and give the encoder your backbuffer surfaces instead.
// Also, asserts? Really? Add some proper error handling already.
// Aaand and and, etc. pp. - but it's a nice starting point to see how HW h.264
// encoding works, and it's tested on AMD, NV and Intel. So have fun :)
#include <assert.h>
#include <stdio.h>
#include <mfidl.h>
#include <mfapi.h>
#include <mferror.h>
#pragma comment(lib, "mfplat.lib")
#pragma comment(lib, "mfuuid.lib")
#define WIDTH 1280
#define HEIGHT 720
#define FPS 30
#define FILENAME L"c:\\temp\\encodetest.264" // play with MPC-HC :)
// render something into an NV12 image buffer
void RenderImage(BYTE *mem, int w, int h)
{
// clear Y with dark grey
memset(mem, 0x40, w*h);
// clear U,V
memset(mem + w*h, 0x80, w*h / 2);
// draw a bright square
static int x = 0;
BYTE *rptr = mem + 100 * w + x;
for (int yy = 0; yy < 50; yy++)
for (int xx = 0; xx < 50; xx++)
rptr[w*yy + xx] = 0xf0;
// .. and animate it
x += 10; if (x + 50 > w) x = 0;
}
// IMFAsyncCallback implementation
struct EncoderCallbacks : IMFAsyncCallback
{
EncoderCallbacks(IMFTransform *encoder)
{
TickEvent = CreateEvent(0, FALSE, FALSE, 0);
encoder->QueryInterface(IID_PPV_ARGS(&Gen));
assert(Gen);
Gen->BeginGetEvent(this, 0);
}
~EncoderCallbacks()
{
Gen->Release();
CloseHandle(TickEvent);
}
// dummy IUnknown impl
virtual HRESULT STDMETHODCALLTYPE QueryInterface(REFIID riid, void ** ppvObject) override { return E_NOTIMPL; }
virtual ULONG STDMETHODCALLTYPE AddRef(void) override { return 1; }
virtual ULONG STDMETHODCALLTYPE Release(void) override { return 1; }
virtual HRESULT STDMETHODCALLTYPE GetParameters(DWORD * pdwFlags, DWORD * pdwQueue) override
{
// we return immediately and don't do anything except signaling another thread
*pdwFlags = MFASYNC_SIGNAL_CALLBACK;
*pdwQueue = MFASYNC_CALLBACK_QUEUE_IO;
return S_OK;
}
virtual HRESULT STDMETHODCALLTYPE Invoke(IMFAsyncResult * pAsyncResult) override
{
IMFMediaEvent *event = 0;
Gen->EndGetEvent(pAsyncResult, &event);
if (event)
{
MediaEventType type;
event->GetType(&type);
switch (type)
{
case METransformNeedInput: InterlockedIncrement(&NeedsInput); break;
case METransformHaveOutput: InterlockedIncrement(&HasOutput); break;
}
event->Release();
SetEvent(TickEvent);
}
Gen->BeginGetEvent(this, 0);
return S_OK;
}
IMFMediaEventGenerator *Gen = nullptr;
HANDLE TickEvent;
unsigned int NeedsInput = 0;
unsigned int HasOutput = 0;
};
int main()
{
HRESULT hr;
// initialize MF
hr = MFStartup(MF_VERSION);
assert(SUCCEEDED(hr));
// enumerate all hardware encoders
IMFActivate **ppActivate = NULL;
UINT32 count = 0;
MFT_REGISTER_TYPE_INFO rtinfo = { MFMediaType_Video, MFVideoFormat_H264 };
MFTEnumEx(MFT_CATEGORY_VIDEO_ENCODER, MFT_ENUM_FLAG_HARDWARE, NULL, &rtinfo, &ppActivate, &count);
assert(ppActivate);
// take first encoder that actually works
IMFTransform *encoder = nullptr;
for (UINT32 i = 0; i < count; i++)
{
if (!encoder)
ppActivate[i]->ActivateObject(IID_PPV_ARGS(&encoder));
ppActivate[i]->Release();
}
assert(encoder);
CoTaskMemFree(ppActivate);
// unlock async mode of encoder
IMFAttributes *encattr = nullptr;
UINT32 async = 0;
encoder->GetAttributes(&encattr);
assert(encattr);
encattr->GetUINT32(MF_TRANSFORM_ASYNC, &async);
assert(async);
encattr->SetUINT32(MF_TRANSFORM_ASYNC_UNLOCK, 1);
// check stream counts just for safety
DWORD inCount = 0, outCount = 0;
encoder->GetStreamCount(&inCount, &outCount);
assert(inCount == 1 && outCount == 1);
// get stream ids
DWORD inStrmId = 0, outStrmId = 0;
encoder->GetStreamIDs(1, &inStrmId, 1, &outStrmId);
if (FAILED(hr)) // can happen, stream IDs are 0 then
inStrmId = outStrmId = 0;
// set media types (hardcoded, but first H264 out and then NV12 in works everywhere)
IMFMediaType *outType = nullptr;
MFCreateMediaType(&outType);
outType->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video);
outType->SetGUID(MF_MT_SUBTYPE, MFVideoFormat_H264);
outType->SetUINT32(MF_MT_COMPRESSED, 1);
outType->SetUINT32(MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive);
MFSetAttributeSize(outType, MF_MT_FRAME_SIZE, WIDTH, HEIGHT);
MFSetAttributeRatio(outType, MF_MT_FRAME_RATE, FPS, 1);
MFSetAttributeRatio(outType, MF_MT_PIXEL_ASPECT_RATIO, 1, 1);
hr = encoder->SetOutputType(outStrmId, outType, 0);
assert(SUCCEEDED(hr));
IMFMediaType *inType = nullptr;
MFCreateMediaType(&inType);
inType->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video);
inType->SetGUID(MF_MT_SUBTYPE, MFVideoFormat_NV12);
inType->SetUINT32(MF_MT_COMPRESSED, 0);
inType->SetUINT32(MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive);
MFSetAttributeSize(inType, MF_MT_FRAME_SIZE, WIDTH, HEIGHT);
MFSetAttributeRatio(inType, MF_MT_FRAME_RATE, FPS, 1);
MFSetAttributeRatio(inType, MF_MT_PIXEL_ASPECT_RATIO, 1, 1);
hr = encoder->SetInputType(inStrmId, inType, 0);
assert(SUCCEEDED(hr));
// just to be safe that the encoder handles samples the way we need
MFT_OUTPUT_STREAM_INFO outInfo;
encoder->GetOutputStreamInfo(outStrmId, &outInfo);
assert(outInfo.dwFlags & (MFT_OUTPUT_STREAM_PROVIDES_SAMPLES | MFT_OUTPUT_STREAM_CAN_PROVIDE_SAMPLES));
// on your marks...
encoder->ProcessMessage(MFT_MESSAGE_NOTIFY_BEGIN_STREAMING, 0);
auto callbacks = new EncoderCallbacks(encoder);
int todo = 500; // # of frames to encode
long long duration = 10ll * 1000ll * 1000ll / ((long long)FPS); // frame duration in 100ns units
long long timecode = 0;
#ifdef FILENAME
HANDLE outFile = CreateFile(FILENAME, GENERIC_WRITE, 0, 0, CREATE_ALWAYS, 0, 0);
assert(outFile != INVALID_HANDLE_VALUE);
#endif
// get set...
encoder->ProcessMessage(MFT_MESSAGE_NOTIFY_START_OF_STREAM, 0);
// go!
while (todo)
{
// wait for any callback to come in
WaitForSingleObject(callbacks->TickEvent, INFINITE);
// as long as the encoder wants new input from us...
while (callbacks->NeedsInput)
{
unsigned int ni = InterlockedDecrement(&callbacks->NeedsInput);
printf("send frame @ %.3fs\n", timecode/10000000.0);
// let's render a frame of video into an MF buffer
IMFMediaBuffer *buffer = nullptr;
MFCreateMemoryBuffer(WIDTH * (HEIGHT + HEIGHT/2), &buffer);
BYTE *memory = nullptr;
buffer->Lock(&memory, 0, 0);
RenderImage(memory, WIDTH, HEIGHT);
buffer->Unlock();
// ... put it into a sample...
IMFSample *sample;
MFCreateSample(&sample);
sample->AddBuffer(buffer);
sample->SetSampleDuration(duration);
sample->SetSampleTime(timecode);
buffer->Release(); // the sample has the buffer now, we don't need it anymore
timecode += duration;
// ... and submit it to the encoder.
hr = encoder->ProcessInput(inStrmId, sample, 0);
assert(SUCCEEDED(hr));
sample->Release(); // the encoder has the sample now, we don't need it anymore
}
// as long as the encoder has stuff for us to process...
while (callbacks->HasOutput)
{
InterlockedDecrement(&callbacks->HasOutput);
DWORD status = 0;
MFT_OUTPUT_DATA_BUFFER outdata = { 0 };
outdata.dwStreamID = outStrmId;
// get sample from encoder
hr = encoder->ProcessOutput(0, 1, &outdata, &status);
if (SUCCEEDED(hr))
{
// for all buffers in the sample...
DWORD bcount;
outdata.pSample->GetBufferCount(&bcount);
for (DWORD i = 0; i < bcount; i++)
{
IMFMediaBuffer *buffer = 0;
outdata.pSample->GetBufferByIndex(i, &buffer);
// get data from buffer
BYTE *data;
DWORD length;
buffer->Lock(&data, 0, &length);
printf("got %d bytes\n", length);
#ifdef FILENAME
DWORD written;
WriteFile(outFile, data, length, &written, 0);
assert(written == length); // seriously now
#endif
buffer->Unlock();
buffer->Release();
}
}
else
{
// some encoders want to renegotiate the output format.
if (hr == MF_E_TRANSFORM_STREAM_CHANGE && (outdata.dwStatus & MFT_OUTPUT_DATA_BUFFER_FORMAT_CHANGE))
{
// Let them know we don't do stuff like that.
hr = encoder->SetOutputType(outStrmId, outType, 0);
assert(SUCCEEDED(hr));
}
else // some other error
assert(0);
}
if (outdata.pSample) outdata.pSample->Release();
if (outdata.pEvents) outdata.pEvents->Release();
if (todo > 0) todo--;
}
}
// we're done here...
encoder->ProcessMessage(MFT_MESSAGE_NOTIFY_END_OF_STREAM, 0);
encoder->ProcessMessage(MFT_MESSAGE_NOTIFY_END_STREAMING, 0);
#ifdef FILENAME
CloseHandle(outFile);
#endif
// shut down
IMFShutdown *shutdown = 0;
encoder->QueryInterface(&shutdown);
if (shutdown)
{
shutdown->Shutdown();
shutdown->Release();
}
delete callbacks;
inType->Release();
outType->Release();
encoder->Release();
encattr->Release();
MFShutdown();
return 0;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment