|
|
@@ -0,0 +1,226 @@
|
|
|
+#include "config.h"
|
|
|
+#include "cuda_helper.hpp"
|
|
|
+#include "video_encoder.h"
|
|
|
+
|
|
|
+#include <nvEncodeAPI.h>
|
|
|
+
|
|
|
+bool check_nvenc_api_call(NVENCSTATUS api_ret, unsigned int line_number,
|
|
|
+ const char *file_name, const char *api_call_str) {
|
|
|
+ if (api_ret == NV_ENC_SUCCESS) [[likely]] return true;
|
|
|
+ SPDLOG_ERROR("NvEnc api call {} failed at {}:{} with error 0x{:x}.",
|
|
|
+ api_call_str, file_name, line_number, api_ret);
|
|
|
+ RET_ERROR;
|
|
|
+}
|
|
|
+
|
|
|
+#define NVENC_API_CHECK(api_call) \
|
|
|
+ if (!check_nvenc_api_call( \
|
|
|
+ api_call, __LINE__, __FILE__, #api_call)) [[unlikely]] \
|
|
|
+ return false
|
|
|
+
|
|
|
+struct video_encoder::impl {
|
|
|
+
|
|
|
+ static constexpr auto frame_buffer_type = NV_ENC_BUFFER_FORMAT_ARGB;
|
|
|
+
|
|
|
+ NV_ENCODE_API_FUNCTION_LIST api = {NV_ENCODE_API_FUNCTION_LIST_VER};
|
|
|
+
|
|
|
+ CUcontext cuda_ctx = nullptr;
|
|
|
+ void *encoder = nullptr;
|
|
|
+ NV_ENC_OUTPUT_PTR output_buf = nullptr;
|
|
|
+
|
|
|
+ int frame_width = image_width * 2, frame_height = image_height;
|
|
|
+ int frame_pitch = frame_width * 4; // ARGB image
|
|
|
+ int frame_rate = default_camera_fps;
|
|
|
+
|
|
|
+ // frame related
|
|
|
+ void *frame_ptr = nullptr, **output_ptr = nullptr;
|
|
|
+ size_t *output_size = nullptr;
|
|
|
+ NV_ENC_REGISTERED_PTR frame_reg_ptr = nullptr;
|
|
|
+
|
|
|
+ bool initialize() {
|
|
|
+ NVENC_API_CHECK(NvEncodeAPICreateInstance(&api));
|
|
|
+ CUDA_API_CHECK(cuCtxGetCurrent(&cuda_ctx));
|
|
|
+ return true;
|
|
|
+ }
|
|
|
+
|
|
|
+ bool start_encode() {
|
|
|
+ // constant params
|
|
|
+ auto codec_guid = NV_ENC_CODEC_HEVC_GUID;
|
|
|
+ auto preset_guid = NV_ENC_PRESET_P3_GUID;
|
|
|
+ auto tuning_info = NV_ENC_TUNING_INFO_LOW_LATENCY;
|
|
|
+
|
|
|
+ // create encoder
|
|
|
+ NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS session_params = {NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS_VER};
|
|
|
+ session_params.deviceType = NV_ENC_DEVICE_TYPE_CUDA;
|
|
|
+ session_params.device = cuda_ctx;
|
|
|
+ session_params.apiVersion = NVENCAPI_VERSION;
|
|
|
+ NVENC_API_CHECK(api.nvEncOpenEncodeSessionEx(&session_params, &encoder));
|
|
|
+
|
|
|
+ // get preset config
|
|
|
+ NV_ENC_PRESET_CONFIG preset_config = {NV_ENC_PRESET_CONFIG_VER, {NV_ENC_CONFIG_VER}};
|
|
|
+ NVENC_API_CHECK(api.nvEncGetEncodePresetConfigEx(
|
|
|
+ encoder, codec_guid, preset_guid, tuning_info, &preset_config));
|
|
|
+ auto &encode_config = preset_config.presetCfg;
|
|
|
+ encode_config.gopLength = NVENC_INFINITE_GOPLENGTH;
|
|
|
+ encode_config.frameIntervalP = 1;
|
|
|
+ auto &rc_params = encode_config.rcParams;
|
|
|
+ rc_params.rateControlMode = NV_ENC_PARAMS_RC_CBR;
|
|
|
+ rc_params.averageBitRate = default_video_stream_bitrate;
|
|
|
+ rc_params.enableAQ = true;
|
|
|
+ rc_params.multiPass = NV_ENC_TWO_PASS_QUARTER_RESOLUTION;
|
|
|
+ // TODO; fine tune encoder config
|
|
|
+
|
|
|
+ // start_encode encoder
|
|
|
+ NV_ENC_INITIALIZE_PARAMS init_params = {NV_ENC_INITIALIZE_PARAMS_VER};
|
|
|
+ init_params.encodeGUID = codec_guid;
|
|
|
+ init_params.presetGUID = preset_guid;
|
|
|
+ init_params.encodeWidth = frame_width;
|
|
|
+ init_params.encodeHeight = frame_height;
|
|
|
+ init_params.darWidth = frame_width; // TODO; learn more about this
|
|
|
+ init_params.darHeight = frame_height; // TODO; learn more about this
|
|
|
+ init_params.frameRateNum = frame_rate;
|
|
|
+ init_params.frameRateDen = 1;
|
|
|
+ init_params.enablePTD = 1;
|
|
|
+ init_params.encodeConfig = &preset_config.presetCfg;
|
|
|
+ init_params.maxEncodeWidth = frame_width;
|
|
|
+ init_params.maxEncodeHeight = frame_height;
|
|
|
+ init_params.tuningInfo = tuning_info;
|
|
|
+ init_params.bufferFormat = frame_buffer_type;
|
|
|
+ NVENC_API_CHECK(api.nvEncInitializeEncoder(encoder, &init_params));
|
|
|
+
|
|
|
+ // create output buffer
|
|
|
+ NV_ENC_CREATE_BITSTREAM_BUFFER buffer_config = {NV_ENC_CREATE_BITSTREAM_BUFFER_VER};
|
|
|
+ NVENC_API_CHECK(api.nvEncCreateBitstreamBuffer(encoder, &buffer_config));
|
|
|
+ output_buf = buffer_config.bitstreamBuffer;
|
|
|
+
|
|
|
+ return true;
|
|
|
+ }
|
|
|
+
|
|
|
+ void cleanup() {
|
|
|
+ if (encoder == nullptr) return;
|
|
|
+
|
|
|
+ // notify the end of stream
|
|
|
+ NV_ENC_PIC_PARAMS pic_params = {NV_ENC_PIC_PARAMS_VER};
|
|
|
+ pic_params.encodePicFlags = NV_ENC_PIC_FLAG_EOS;
|
|
|
+ api.nvEncEncodePicture(encoder, &pic_params);
|
|
|
+
|
|
|
+ // releasing resources
|
|
|
+ if (frame_ptr != nullptr) {
|
|
|
+ unregister_frame_ptr();
|
|
|
+ frame_ptr = nullptr;
|
|
|
+ }
|
|
|
+ api.nvEncDestroyBitstreamBuffer(encoder, output_buf);
|
|
|
+
|
|
|
+ // close encoder
|
|
|
+ api.nvEncDestroyEncoder(encoder);
|
|
|
+ encoder = nullptr;
|
|
|
+ }
|
|
|
+
|
|
|
+ bool register_frame_ptr() {
|
|
|
+ NV_ENC_REGISTER_RESOURCE reg_params = {NV_ENC_REGISTER_RESOURCE_VER};
|
|
|
+ reg_params.resourceType = NV_ENC_INPUT_RESOURCE_TYPE_CUDADEVICEPTR;
|
|
|
+ reg_params.width = frame_width;
|
|
|
+ reg_params.height = frame_height;
|
|
|
+ reg_params.pitch = frame_pitch;
|
|
|
+ reg_params.resourceToRegister = (void *) frame_ptr;
|
|
|
+ reg_params.bufferFormat = frame_buffer_type;
|
|
|
+ reg_params.bufferUsage = NV_ENC_INPUT_IMAGE;
|
|
|
+ NVENC_API_CHECK(api.nvEncRegisterResource(encoder, ®_params));
|
|
|
+ frame_reg_ptr = reg_params.registeredResource;
|
|
|
+ return true;
|
|
|
+ }
|
|
|
+
|
|
|
+ bool unregister_frame_ptr() {
|
|
|
+ if (frame_reg_ptr == nullptr) return true;
|
|
|
+ NVENC_API_CHECK(api.nvEncUnregisterResource(encoder, frame_reg_ptr));
|
|
|
+ return true;
|
|
|
+ }
|
|
|
+
|
|
|
+ bool encode_frame() {
|
|
|
+ // map input resource
|
|
|
+ NV_ENC_MAP_INPUT_RESOURCE map_params = {NV_ENC_MAP_INPUT_RESOURCE_VER};
|
|
|
+ map_params.registeredResource = frame_reg_ptr;
|
|
|
+ NVENC_API_CHECK(api.nvEncMapInputResource(encoder, &map_params));
|
|
|
+ assert(map_params.mappedBufferFmt == frame_buffer_type);
|
|
|
+
|
|
|
+ // encode frame
|
|
|
+ NV_ENC_PIC_PARAMS pic_params = {NV_ENC_PIC_PARAMS_VER};
|
|
|
+ pic_params.inputWidth = frame_width;
|
|
|
+ pic_params.inputHeight = frame_height;
|
|
|
+ pic_params.inputPitch = frame_pitch;
|
|
|
+ pic_params.encodePicFlags = 0; // TODO; modify this value
|
|
|
+ pic_params.inputBuffer = map_params.mappedResource;
|
|
|
+ pic_params.outputBitstream = output_buf;
|
|
|
+ pic_params.bufferFmt = frame_buffer_type;
|
|
|
+ pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FRAME; // TODO; learn more about this
|
|
|
+ NVENC_API_CHECK(api.nvEncEncodePicture(encoder, &pic_params));
|
|
|
+
|
|
|
+ // get encoded bitstream
|
|
|
+ NV_ENC_LOCK_BITSTREAM lock_config = {NV_ENC_LOCK_BITSTREAM_VER};
|
|
|
+ lock_config.doNotWait = false; // block until encode completed.
|
|
|
+ lock_config.outputBitstream = output_buf;
|
|
|
+ NVENC_API_CHECK(api.nvEncLockBitstream(encoder, &lock_config));
|
|
|
+
|
|
|
+ // copy bitstream
|
|
|
+ assert(output_ptr != nullptr);
|
|
|
+ assert(output_size != nullptr);
|
|
|
+ *output_ptr = malloc(lock_config.bitstreamSizeInBytes);
|
|
|
+ *output_size = lock_config.bitstreamSizeInBytes;
|
|
|
+ memcpy(*output_ptr, lock_config.bitstreamBufferPtr, lock_config.bitstreamSizeInBytes);
|
|
|
+
|
|
|
+ // cleanup
|
|
|
+ NVENC_API_CHECK(api.nvEncUnlockBitstream(encoder, output_buf));
|
|
|
+ NVENC_API_CHECK(api.nvEncUnmapInputResource(encoder, map_params.mappedResource));
|
|
|
+
|
|
|
+ return true;
|
|
|
+ }
|
|
|
+
|
|
|
+};
|
|
|
+
|
|
|
+video_encoder::video_encoder()
|
|
|
+ : pimpl(std::make_unique<impl>()) {}
|
|
|
+
|
|
|
+video_encoder::~video_encoder() = default;
|
|
|
+
|
|
|
+bool video_encoder::initialize() {
|
|
|
+ return pimpl->initialize();
|
|
|
+}
|
|
|
+
|
|
|
+bool video_encoder::start_encode(int width, int height, int fps) {
|
|
|
+ pimpl->frame_width = width;
|
|
|
+ pimpl->frame_height = height;
|
|
|
+ pimpl->frame_pitch = width * 4; // ARGB image
|
|
|
+ pimpl->frame_rate = fps;
|
|
|
+ return pimpl->start_encode();
|
|
|
+}
|
|
|
+
|
|
|
+void video_encoder::stop_encode() {
|
|
|
+ pimpl->cleanup();
|
|
|
+}
|
|
|
+
|
|
|
+bool video_encoder::encode_frame(void *frame_ptr, void **output_ptr, size_t *output_size) {
|
|
|
+ // register frame ptr
|
|
|
+ if (pimpl->frame_ptr != frame_ptr) {
|
|
|
+ pimpl->unregister_frame_ptr();
|
|
|
+ pimpl->frame_ptr = frame_ptr;
|
|
|
+ pimpl->register_frame_ptr();
|
|
|
+ }
|
|
|
+
|
|
|
+ pimpl->output_ptr = output_ptr;
|
|
|
+ pimpl->output_size = output_size;
|
|
|
+ return pimpl->encode_frame();
|
|
|
+}
|
|
|
+
|
|
|
+bool video_encoder::encode_frame(cudaGraphicsResource *res, void **output_ptr, size_t *output_size) {
|
|
|
+ void *pbo_ptr;
|
|
|
+ size_t pbo_size;
|
|
|
+ CUDA_API_CHECK(cudaGraphicsMapResources(1, &res));
|
|
|
+ CUDA_API_CHECK(cudaGraphicsResourceGetMappedPointer(&pbo_ptr, &pbo_size, res));
|
|
|
+ assert(pbo_size == pimpl->frame_pitch * pimpl->frame_height);
|
|
|
+ CALL_CHECK(encode_frame(pbo_ptr, output_ptr, output_size));
|
|
|
+ CUDA_API_CHECK(cudaGraphicsUnmapResources(1, &res));
|
|
|
+ return true;
|
|
|
+}
|
|
|
+
|
|
|
+bool video_encoder::is_encoding() {
|
|
|
+ return pimpl->encoder != nullptr;
|
|
|
+}
|