|
@@ -0,0 +1,266 @@
|
|
|
|
|
+#include "third_party/scope_guard.hpp"
|
|
|
|
|
+#include "encoder_nvenc.h"
|
|
|
|
|
+
|
|
|
|
|
+#include <nvEncodeAPI.h>
|
|
|
|
|
+
|
|
|
|
|
+#include <fmt/chrono.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, (int) api_ret);
|
|
|
|
|
+ RET_ERROR_B;
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+#define API_CHECK(api_call) \
|
|
|
|
|
+ check_nvenc_api_call( \
|
|
|
|
|
+ api_call, __LINE__, __FILE__, #api_call)
|
|
|
|
|
+
|
|
|
|
|
+#define API_CHECK_P(api_call) \
|
|
|
|
|
+ if (!check_nvenc_api_call( \
|
|
|
|
|
+ api_call, __LINE__, __FILE__, #api_call)) [[unlikely]] \
|
|
|
|
|
+ return nullptr
|
|
|
|
|
+
|
|
|
|
|
+namespace video_encoder_impl {
|
|
|
|
|
+ constexpr auto frame_buffer_type = NV_ENC_BUFFER_FORMAT_ARGB;
|
|
|
|
|
+ static auto codec_guid = NV_ENC_CODEC_HEVC_GUID;
|
|
|
|
|
+ static auto preset_guid = NV_ENC_PRESET_P3_GUID;
|
|
|
|
|
+ constexpr auto tuning_info = NV_ENC_TUNING_INFO_ULTRA_LOW_LATENCY;
|
|
|
|
|
+ std::unique_ptr<NV_ENCODE_API_FUNCTION_LIST> api;
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+using namespace video_encoder_impl;
|
|
|
|
|
+
|
|
|
|
|
+struct encoder_nvenc::impl {
|
|
|
|
|
+
|
|
|
|
|
+ std::unique_ptr<NV_ENC_PRESET_CONFIG> preset_config;
|
|
|
|
|
+ std::unique_ptr<NV_ENC_INITIALIZE_PARAMS> init_params;
|
|
|
|
|
+ void *encoder;
|
|
|
|
|
+ NV_ENC_OUTPUT_PTR output_buf;
|
|
|
|
|
+
|
|
|
|
|
+ cv::Size frame_size;
|
|
|
|
|
+ FILE *save_file = nullptr;
|
|
|
|
|
+ bool save_length;
|
|
|
|
|
+
|
|
|
|
|
+ void *last_frame_ptr = nullptr;
|
|
|
|
|
+ NV_ENC_REGISTERED_PTR last_reg_ptr = nullptr;
|
|
|
|
|
+ uint64_t last_frame_id = 0;
|
|
|
|
|
+
|
|
|
|
|
+ smart_cuda_stream *stream = nullptr;
|
|
|
|
|
+
|
|
|
|
|
+ ~impl() {
|
|
|
|
|
+ // 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_CHECK(api->nvEncEncodePicture(encoder, &pic_params));
|
|
|
|
|
+
|
|
|
|
|
+ // releasing resources
|
|
|
|
|
+ unregister_frame_ptr();
|
|
|
|
|
+ API_CHECK(api->nvEncDestroyBitstreamBuffer(encoder, output_buf));
|
|
|
|
|
+
|
|
|
|
|
+ // close encoder
|
|
|
|
|
+ API_CHECK(api->nvEncDestroyEncoder(encoder));
|
|
|
|
|
+
|
|
|
|
|
+ // close save file
|
|
|
|
|
+ if (save_file != nullptr) {
|
|
|
|
|
+ fclose(save_file);
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ SPDLOG_INFO("Video encoder stopped.");
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ static impl *create(create_config conf) {
|
|
|
|
|
+ // initialize api
|
|
|
|
|
+ if (api == nullptr) [[unlikely]] {
|
|
|
|
|
+ api = std::make_unique<NV_ENCODE_API_FUNCTION_LIST>(
|
|
|
|
|
+ NV_ENCODE_API_FUNCTION_LIST_VER);
|
|
|
|
|
+ API_CHECK_P(NvEncodeAPICreateInstance(api.get()));
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ // get cuda context
|
|
|
|
|
+ auto cuda_ctx = conf.ctx;
|
|
|
|
|
+
|
|
|
|
|
+ // create encoder
|
|
|
|
|
+ auto ret = new impl;
|
|
|
|
|
+ ret->stream = conf.stream;
|
|
|
|
|
+ ret->frame_size = conf.frame_size;
|
|
|
|
|
+ auto closer = sg::make_scope_guard([&] {
|
|
|
|
|
+ if (ret->save_file != nullptr) {
|
|
|
|
|
+ fclose(ret->save_file);
|
|
|
|
|
+ }
|
|
|
|
|
+ delete ret;
|
|
|
|
|
+ });
|
|
|
|
|
+ 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;
|
|
|
|
|
+ API_CHECK_P(api->nvEncOpenEncodeSessionEx(&session_params, &ret->encoder));
|
|
|
|
|
+
|
|
|
|
|
+ // get preset config
|
|
|
|
|
+ ret->preset_config = std::make_unique<NV_ENC_PRESET_CONFIG>();
|
|
|
|
|
+ auto &preset_config = *ret->preset_config;
|
|
|
|
|
+ preset_config.version = NV_ENC_PRESET_CONFIG_VER;
|
|
|
|
|
+ preset_config.presetCfg.version = NV_ENC_CONFIG_VER;
|
|
|
|
|
+ API_CHECK_P(api->nvEncGetEncodePresetConfigEx(
|
|
|
|
|
+ ret->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 = conf.bitrate_mbps * 1e6;
|
|
|
|
|
+ rc_params.enableAQ = true;
|
|
|
|
|
+ rc_params.multiPass = NV_ENC_TWO_PASS_QUARTER_RESOLUTION;
|
|
|
|
|
+ // TODO; fine tune encoder config
|
|
|
|
|
+
|
|
|
|
|
+ // start_encode encoder
|
|
|
|
|
+ ret->init_params =
|
|
|
|
|
+ std::make_unique<NV_ENC_INITIALIZE_PARAMS>(NV_ENC_INITIALIZE_PARAMS_VER);
|
|
|
|
|
+ auto &init_params = *ret->init_params;
|
|
|
|
|
+ init_params.encodeGUID = codec_guid;
|
|
|
|
|
+ init_params.presetGUID = preset_guid;
|
|
|
|
|
+ init_params.encodeWidth = conf.frame_size.width;
|
|
|
|
|
+ init_params.encodeHeight = conf.frame_size.height;
|
|
|
|
|
+ init_params.darWidth = conf.frame_size.width; // TODO; learn more about this
|
|
|
|
|
+ init_params.darHeight = conf.frame_size.height; // TODO; learn more about this
|
|
|
|
|
+ init_params.frameRateNum = conf.frame_rate;
|
|
|
|
|
+ init_params.frameRateDen = 1;
|
|
|
|
|
+ init_params.enablePTD = 1;
|
|
|
|
|
+ init_params.encodeConfig = &preset_config.presetCfg;
|
|
|
|
|
+ init_params.maxEncodeWidth = conf.frame_size.width;
|
|
|
|
|
+ init_params.maxEncodeHeight = conf.frame_size.height;
|
|
|
|
|
+ init_params.tuningInfo = tuning_info;
|
|
|
|
|
+ init_params.bufferFormat = frame_buffer_type;
|
|
|
|
|
+ API_CHECK_P(api->nvEncInitializeEncoder(ret->encoder, &init_params));
|
|
|
|
|
+
|
|
|
|
|
+ // create output buffer
|
|
|
|
|
+ NV_ENC_CREATE_BITSTREAM_BUFFER buffer_config = {
|
|
|
|
|
+ NV_ENC_CREATE_BITSTREAM_BUFFER_VER};
|
|
|
|
|
+ API_CHECK_P(api->nvEncCreateBitstreamBuffer(ret->encoder, &buffer_config));
|
|
|
|
|
+ ret->output_buf = buffer_config.bitstreamBuffer;
|
|
|
|
|
+
|
|
|
|
|
+ // create save file
|
|
|
|
|
+ if (conf.save_file) {
|
|
|
|
|
+ auto file_name = fmt::format("record_{:%Y_%m_%d_%H_%M_%S}.{}",
|
|
|
|
|
+ now_since_epoch_in_seconds(),
|
|
|
|
|
+ conf.save_length ? "dat" : "hevc");
|
|
|
|
|
+ ret->save_file = fopen(file_name.c_str(), "wb");
|
|
|
|
|
+ ret->save_length = conf.save_length;
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ SPDLOG_INFO("Video encoder started.");
|
|
|
|
|
+ closer.dismiss();
|
|
|
|
|
+ return ret;
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ void unregister_frame_ptr() {
|
|
|
|
|
+ if (last_reg_ptr == nullptr) return;
|
|
|
|
|
+ API_CHECK(api->nvEncUnregisterResource(encoder, last_reg_ptr));
|
|
|
|
|
+ last_reg_ptr = nullptr;
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ void register_frame_ptr(const image_info_type<uchar4> &info) {
|
|
|
|
|
+ assert(info.loc == MEM_CUDA);
|
|
|
|
|
+ NV_ENC_REGISTER_RESOURCE reg_params = {NV_ENC_REGISTER_RESOURCE_VER};
|
|
|
|
|
+ reg_params.resourceType = NV_ENC_INPUT_RESOURCE_TYPE_CUDADEVICEPTR;
|
|
|
|
|
+ reg_params.width = info.size.width;
|
|
|
|
|
+ reg_params.height = info.size.height;
|
|
|
|
|
+ reg_params.pitch = info.pitch;
|
|
|
|
|
+ reg_params.resourceToRegister = info.start_ptr();
|
|
|
|
|
+ reg_params.bufferFormat = frame_buffer_type;
|
|
|
|
|
+ reg_params.bufferUsage = NV_ENC_INPUT_IMAGE;
|
|
|
|
|
+ API_CHECK(api->nvEncRegisterResource(encoder, ®_params));
|
|
|
|
|
+ last_reg_ptr = reg_params.registeredResource;
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ frame_info encode(const image_u8c4 &img, bool force_idr = false) {
|
|
|
|
|
+ // register pointer if needed
|
|
|
|
|
+ auto img_info = img->as_cuda_info(stream);
|
|
|
|
|
+ // TODO: image pointer may change frequently
|
|
|
|
|
+ if (img_info.start_ptr() != last_frame_ptr) [[unlikely]] {
|
|
|
|
|
+ assert(img->size() == frame_size);
|
|
|
|
|
+ unregister_frame_ptr();
|
|
|
|
|
+ register_frame_ptr(img_info);
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ // map input resource
|
|
|
|
|
+ NV_ENC_MAP_INPUT_RESOURCE map_params = {
|
|
|
|
|
+ NV_ENC_MAP_INPUT_RESOURCE_VER};
|
|
|
|
|
+ map_params.registeredResource = last_reg_ptr;
|
|
|
|
|
+ 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 = img_info.size.width;
|
|
|
|
|
+ pic_params.inputHeight = img_info.size.height;
|
|
|
|
|
+ pic_params.inputPitch = img_info.pitch;
|
|
|
|
|
+ if (force_idr) { // request for IDR frame
|
|
|
|
|
+ pic_params.encodePicFlags = NV_ENC_PIC_FLAG_FORCEIDR | NV_ENC_PIC_FLAG_OUTPUT_SPSPPS;
|
|
|
|
|
+ } else {
|
|
|
|
|
+ pic_params.encodePicFlags = 0;
|
|
|
|
|
+ }
|
|
|
|
|
+ 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
|
|
|
|
|
+ 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;
|
|
|
|
|
+ API_CHECK(api->nvEncLockBitstream(encoder, &lock_config));
|
|
|
|
|
+
|
|
|
|
|
+ // copy bitstream
|
|
|
|
|
+ auto ret = frame_info();
|
|
|
|
|
+ ret.data = data_type(lock_config.bitstreamSizeInBytes, lock_config.bitstreamBufferPtr);
|
|
|
|
|
+ ret.idr = lock_config.pictureType == NV_ENC_PIC_TYPE_IDR;
|
|
|
|
|
+ ret.frame_id = ++last_frame_id;
|
|
|
|
|
+
|
|
|
|
|
+ // save bitstream
|
|
|
|
|
+ if (save_file != nullptr) {
|
|
|
|
|
+ auto ret_size = ret.size();
|
|
|
|
|
+ if (save_length) {
|
|
|
|
|
+ fwrite(&ret_size, sizeof(size_t), 1, save_file);
|
|
|
|
|
+ }
|
|
|
|
|
+ fwrite(ret.start_ptr(), ret_size, 1, save_file);
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ // cleanup
|
|
|
|
|
+ API_CHECK(api->nvEncUnlockBitstream(encoder, output_buf));
|
|
|
|
|
+ API_CHECK(api->nvEncUnmapInputResource(encoder, map_params.mappedResource));
|
|
|
|
|
+
|
|
|
|
|
+ return ret;
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ void change_config(modifiable_config conf) {
|
|
|
|
|
+ NV_ENC_RECONFIGURE_PARAMS params = {NV_ENC_RECONFIGURE_PARAMS_VER};
|
|
|
|
|
+ init_params->frameRateNum = conf.frame_rate;
|
|
|
|
|
+ init_params->encodeConfig->rcParams.averageBitRate = conf.bitrate_mbps * 1e6;
|
|
|
|
|
+ params.reInitEncodeParams = *init_params;
|
|
|
|
|
+ params.resetEncoder = true;
|
|
|
|
|
+ params.forceIDR = true;
|
|
|
|
|
+ API_CHECK(api->nvEncReconfigureEncoder(encoder, ¶ms));
|
|
|
|
|
+ }
|
|
|
|
|
+};
|
|
|
|
|
+
|
|
|
|
|
+encoder_nvenc::~encoder_nvenc() = default;
|
|
|
|
|
+
|
|
|
|
|
+encoder_nvenc::pointer encoder_nvenc::create(create_config conf) {
|
|
|
|
|
+ auto pimpl = impl::create(conf);
|
|
|
|
|
+ if (pimpl == nullptr) return nullptr;
|
|
|
|
|
+ auto ret = std::make_unique<encoder_nvenc>();
|
|
|
|
|
+ ret->pimpl.reset(pimpl);
|
|
|
|
|
+ return ret;
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+frame_info encoder_nvenc::encode(const image_u8c4 &img, bool force_idr) {
|
|
|
|
|
+ return pimpl->encode(img, force_idr);
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+void encoder_nvenc::change_config(modifiable_config conf) {
|
|
|
|
|
+ pimpl->change_config(conf);
|
|
|
|
|
+}
|