amf_bgra_nv12_test.cpp

#include <stdio.h>

#include <fstream>
#include <thread>

#include "public/common/AMFFactory.h"
#include "public/common/AMFSTL.h"
#include "public/common/Thread.h"
#include "public/common/TraceAdapter.h"
#include "public/include/components/VideoEncoderVCE.h"

#define CHECK_AMF_RET(AMF_call)                        \
  {                                                    \
    AMF_RESULT r = AMF_call;                           \
    if (r != AMF_OK && r != AMF_ALREADY_INITIALIZED) { \
      printf("fail with error code %d", r);            \
      return false;                                    \
    }                                                  \
  }

// Use BGRA if defined. Otherwise, use NV12.
#define BGRA_MODE

static amf::AMF_MEMORY_TYPE memoryTypeIn = amf::AMF_MEMORY_VULKAN;
#ifdef BGRA_MODE
static amf::AMF_SURFACE_FORMAT formatIn = amf::AMF_SURFACE_BGRA;
#else
static amf::AMF_SURFACE_FORMAT formatIn = amf::AMF_SURFACE_NV12;
#endif

static amf_int32 widthIn = 1920;
static amf_int32 heightIn = 1080;
static amf_int32 frameRateIn = 30;
static amf_int64 bitRateIn = 5000000L;
static amf_int32 rectSize = 50;
static amf_int32 frameCount = 30;

amf::AMFSurfacePtr pColor1;

class PollingThread : public amf::AMFThread {
 public:
  PollingThread(amf::AMFComponent *encoder) : m_pEncoder(encoder) {}
  virtual ~PollingThread() {}
  virtual void Run() {
    RequestStop();
    while (true) {
      amf::AMFDataPtr data;
      AMF_RESULT res = m_pEncoder->QueryOutput(&data);
      if (res == AMF_EOF) {
        break;
      }
      if (data != NULL) {
        // do nothing. just consume data.
      } else {
        amf_sleep(1);
      }
    }
    m_pEncoder = NULL;
  }

 private:
  amf::AMFComponentPtr m_pEncoder;
};

static void FillNV12SurfaceWithColor(amf::AMFSurface *surface, amf_uint8 Y, amf_uint8 U, amf_uint8 V) {
  amf::AMFPlane *pPlaneY = surface->GetPlaneAt(0);
  amf::AMFPlane *pPlaneUV = surface->GetPlaneAt(1);

  amf_int32 lineY = pPlaneY->GetHPitch();
  amf_uint8 *pDataUV = (amf_uint8 *)pPlaneUV->GetNative();

  for (amf_int32 y = 0; y < pPlaneY->GetHeight(); y++) {
    amf_uint8 *pDataLine = (amf_uint8 *)pPlaneY->GetNative() + y * lineY;
    memset(pDataLine, Y, pPlaneY->GetWidth());
  }

  for (amf_int32 y = 0; y < pPlaneUV->GetHeight(); y++) {
    amf_uint8 *pDataLine = (amf_uint8 *)pDataUV + y * pPlaneUV->GetHPitch();
    for (amf_int32 x = 0; x < pPlaneUV->GetWidth(); x++) {
      *pDataLine++ = U;
      *pDataLine++ = V;
    }
  }
}

static void FillBGRASurfaceWithColor(amf::AMFSurface *surface, amf_uint8 B, amf_uint8 G, amf_uint8 R, amf_uint8 A) {
  amf::AMFPlane *plane = surface->GetPlane(amf::AMF_PLANE_PACKED);
  amf_uint8 *data = (amf_uint8 *)plane->GetNative();

  for (amf_int32 y = 0; y < plane->GetHeight(); y++) {
    amf_uint8 *line = (amf_uint8 *)data + y * plane->GetHPitch();
    for (amf_int32 x = 0; x < plane->GetWidth(); x++) {
      *line++ = B;
      *line++ = G;
      *line++ = R;
      *line++ = A;
    }
  }
}

static void PrepareFillFromHost(amf::AMFContext *context) {
  AMF_RESULT res = AMF_OK;
  res = context->AllocSurface(amf::AMF_MEMORY_HOST, formatIn, widthIn, heightIn, &pColor1);
#ifdef BGRA_MODE
  FillBGRASurfaceWithColor(pColor1, 128, 255, 128, 255);
#else
  FillNV12SurfaceWithColor(pColor1, 128, 255, 128);
#endif
  pColor1->Convert(memoryTypeIn);
}

static void FillSurfaceVulkan(amf::AMFContext *context, amf::AMFSurface *surface) {
  amf::AMFComputePtr compute;
  context->GetCompute(amf::AMF_MEMORY_VULKAN, &compute);
#ifdef BGRA_MODE
  amf::AMFPlane *plane = pColor1->GetPlane(amf::AMF_PLANE_PACKED);
  amf_size origin1[3] = {0, 0, 0};
  amf_size region1[3] = {(amf_size)plane->GetWidth(), (amf_size)plane->GetHeight(), (amf_size)1};
  compute->CopyPlane(plane, origin1, region1, surface->GetPlane(amf::AMF_PLANE_PACKED), origin1);
#else
  for (int p = 0; p < 2; p++) {
    amf::AMFPlane *plane = pColor1->GetPlaneAt(p);
    amf_size origin1[3] = {0, 0, 0};
    amf_size region1[3] = {(amf_size)plane->GetWidth(), (amf_size)plane->GetHeight(), (amf_size)1};
    compute->CopyPlane(plane, origin1, region1, surface->GetPlaneAt(p), origin1);
  }
#endif
}

int main(int argc, char *argv[]) {
  AMF_RESULT res = AMF_OK;  // error checking can be added later
  res = g_AMFFactory.Init();
  if (res != AMF_OK) {
    wprintf(L"AMF Failed to initialize");
    return 1;
  }

  amf::AMFTraceEnableWriter(AMF_TRACE_WRITER_CONSOLE, true);
  amf::AMFTraceEnableWriter(AMF_TRACE_WRITER_DEBUG_OUTPUT, true);

  amf::AMFContextPtr context;
  amf::AMFComponentPtr encoder;
  amf::AMFSurfacePtr surfaceIn;

  CHECK_AMF_RET(g_AMFFactory.GetFactory()->CreateContext(&context));
  CHECK_AMF_RET(amf::AMFContext1Ptr(context)->InitVulkan(NULL));
  PrepareFillFromHost(context);

  CHECK_AMF_RET(g_AMFFactory.GetFactory()->CreateComponent(context, AMFVideoEncoderVCE_AVC, &encoder));
  CHECK_AMF_RET(encoder->SetProperty(AMF_VIDEO_ENCODER_USAGE, AMF_VIDEO_ENCODER_USAGE_TRANSCONDING));
  CHECK_AMF_RET(encoder->SetProperty(AMF_VIDEO_ENCODER_B_PIC_PATTERN, 0));
  CHECK_AMF_RET(encoder->SetProperty(AMF_VIDEO_ENCODER_QUALITY_PRESET, AMF_VIDEO_ENCODER_QUALITY_PRESET_SPEED));
  CHECK_AMF_RET(encoder->SetProperty(AMF_VIDEO_ENCODER_TARGET_BITRATE, bitRateIn));
  CHECK_AMF_RET(encoder->SetProperty(AMF_VIDEO_ENCODER_FRAMESIZE, ::AMFConstructSize(widthIn, heightIn)));
  CHECK_AMF_RET(encoder->SetProperty(AMF_VIDEO_ENCODER_FRAMERATE, ::AMFConstructRate(frameRateIn, 1)));
  CHECK_AMF_RET(encoder->Init(formatIn, widthIn, heightIn));

  PollingThread thread(encoder);
  thread.Start();

  amf_int32 submitted = 0;
  while (submitted < frameCount) {
    // Change resolution after 10 successful submit.
    if (submitted == 10 && res == AMF_OK) {
      printf("resize from %dx%d to %dx%d\n", widthIn, heightIn, widthIn + 100, heightIn);
      widthIn += 100;
      CHECK_AMF_RET(encoder->ReInit(widthIn, heightIn));
    }

    if (surfaceIn == NULL) {
      CHECK_AMF_RET(context->AllocSurface(memoryTypeIn, formatIn, widthIn, heightIn, &surfaceIn));
      FillSurfaceVulkan(context, surfaceIn);
    }

    res = encoder->SubmitInput(surfaceIn);
    if (res == AMF_INPUT_FULL) {
      amf_sleep(10);
    } else {
      printf("%d result code: %d\n", submitted, res);
      surfaceIn = NULL;
      submitted++;
    }
  }
  while (true) {
    res = encoder->Drain();
    if (res != AMF_INPUT_FULL) {
      break;
    }
    amf_sleep(1);
  }
  thread.WaitForStop();

  pColor1 = NULL;
  surfaceIn = NULL;
  encoder->Terminate();
  encoder = NULL;
  context->Terminate();
  context = NULL;

  g_AMFFactory.Terminate();
  return 0;
}