Browse Source

Implemented NvEnc video encoding.

jcsyshc 2 năm trước cách đây
mục cha
commit
518143ef78

+ 6 - 4
CMakeLists.txt

@@ -1,5 +1,5 @@
 cmake_minimum_required(VERSION 3.25)
-project(RemoteAR2 LANGUAGES C CXX CUDA)
+project(RemoteAR2 LANGUAGES C CXX)
 
 set(CMAKE_CXX_STANDARD 20)
 
@@ -52,6 +52,7 @@ target_sources(${PROJECT_NAME} PRIVATE
 # spdlog config
 find_package(spdlog REQUIRED)
 target_link_libraries(${PROJECT_NAME} spdlog::spdlog)
+target_compile_definitions(${PROJECT_NAME} PRIVATE SPDLOG_ACTIVE_LEVEL=SPDLOG_LEVEL_TRACE)
 
 # OpenCV config
 find_package(OpenCV REQUIRED COMPONENTS cudaimgproc imgcodecs)
@@ -101,8 +102,8 @@ find_package(Eigen3 REQUIRED)
 target_link_libraries(${PROJECT_NAME} Eigen3::Eigen)
 
 # CUDA config
-target_include_directories(${PROJECT_NAME} PRIVATE ${CUDA_INCLUDE_DIRS})
-target_link_libraries(${PROJECT_NAME} ${CUDA_LIBRARIES})
+find_package(CUDAToolkit REQUIRED)
+target_link_libraries(${PROJECT_NAME} CUDA::cudart CUDA::cuda_driver)
 
 # NvEnc config
 if (WIN32)
@@ -115,4 +116,5 @@ else ()
 endif ()
 set(NVCODEC_INCLUDE_DIR ${NVCODEC_DIR}/Interface)
 target_include_directories(${PROJECT_NAME} PRIVATE ${NVCODEC_INCLUDE_DIR})
-target_link_libraries(${PROJECT_NAME} ${NVENC_LIB})
+target_link_libraries(${PROJECT_NAME} ${NVENC_LIB})
+target_sources(${PROJECT_NAME} PRIVATE src/video_encoder.cpp)

+ 6 - 2
src/augment_renderer.cpp

@@ -13,6 +13,10 @@ struct augment_renderer::impl {
     texture_renderer *tex_renderer = nullptr;
     const render_config *config = nullptr;
 
+    ~impl() {
+        cudaGraphicsUnregisterResource(bg_res);
+    }
+
     bool initialize() {
         // generate and allocate pixel buffer
         glGenBuffers(1, &bg_pbo);
@@ -91,7 +95,7 @@ void augment_renderer::set_background(const cv::cuda::GpuMat *background) {
     pimpl->bg_img = background;
 }
 
-bool augment_renderer::render(const render_config *config) {
-    pimpl->config = config;
+bool augment_renderer::render(const render_config &config) {
+    pimpl->config = &config;
     return pimpl->render();
 }

+ 1 - 1
src/augment_renderer.h

@@ -23,7 +23,7 @@ public:
         float width, height;
     };
 
-    bool render(const render_config *config);
+    bool render(const render_config &config);
 
 private:
     struct impl;

+ 4 - 1
src/config.h

@@ -20,11 +20,14 @@ static constexpr auto right_camera_name = "RightEye";
 
 static constexpr auto default_camera_exposure_time_ms = 5; // 5ms
 static constexpr auto default_camera_analog_gain = 20; // 20dB
-static constexpr auto default_camera_capture_interval_ms = 33; // 33ms
+static constexpr auto default_camera_fps = 30; // 30 fps
 
 static constexpr auto default_time_out = std::chrono::milliseconds(50); // 50ms
 static constexpr auto default_spin_time = std::chrono::milliseconds(100); // 100us
 
+static constexpr auto default_cuda_device_id = 0;
+static constexpr auto default_video_stream_bitrate = 10 * 1e6; // 10mbps
+
 #define RET_ERROR \
     assert(false); \
     return false; \

+ 10 - 0
src/cuda_helper.hpp

@@ -33,4 +33,14 @@ inline bool check_cuda_api_call(cudaError api_ret, unsigned int line_number,
         api_call, __LINE__, __FILE__, #api_call)) [[unlikely]] \
         return false
 
+inline bool create_cuda_context(CUcontext *ctx) {
+    int cuda_device_count;
+    CUDA_API_CHECK(cuDeviceGetCount(&cuda_device_count));
+    CALL_CHECK(cuda_device_count > default_cuda_device_id);
+    CUdevice cuda_device;
+    CUDA_API_CHECK(cuDeviceGet(&cuda_device, default_cuda_device_id));
+    CUDA_API_CHECK(cuCtxCreate(ctx, CU_CTX_SCHED_AUTO, cuda_device));
+    return true;
+}
+
 #endif //REMOTEAR2_CUDA_HELPER_HPP

+ 74 - 0
src/frame_buffer_helper.hpp

@@ -0,0 +1,74 @@
+#ifndef REMOTEAR2_FRAME_BUFFER_HELPER_HPP
+#define REMOTEAR2_FRAME_BUFFER_HELPER_HPP
+
+#include "config.h"
+#include "cuda_helper.hpp"
+
+#include <cuda_gl_interop.h>
+
+#include <glad/gl.h>
+
+struct frame_buffer_helper {
+
+    int tex_width = 2 * image_width, tex_height = image_height;
+    GLuint tex = 0, fbo = 0, pbo = 0;
+    cudaGraphicsResource *pbo_res = nullptr;
+
+    ~frame_buffer_helper() {
+        if (fbo == 0)return;
+        cudaGraphicsUnregisterResource(pbo_res);
+        glDeleteBuffers(1, &pbo);
+        glDeleteFramebuffers(1, &fbo);
+        glDeleteTextures(1, &tex);
+    }
+
+    bool initialize(int width, int height) {
+        tex_width = width;
+        tex_height = height;
+
+        // create and allocate tex
+        glGenTextures(1, &tex);
+        glBindTexture(GL_TEXTURE_2D, tex);
+        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
+        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
+        glTexStorage2D(GL_TEXTURE_2D, 1, GL_RGB8, tex_width, tex_height);
+
+        // create and config fbo
+        glGenFramebuffers(1, &fbo);
+        glBindFramebuffer(GL_FRAMEBUFFER, fbo);
+        glBindTexture(GL_TEXTURE_2D, tex);
+        glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, tex, 0);
+        CALL_CHECK(check_frame_buffer());
+
+        // create and config pbo
+        glGenBuffers(1, &pbo);
+        glBindBuffer(GL_PIXEL_PACK_BUFFER, pbo);
+        glBufferStorage(GL_PIXEL_PACK_BUFFER, tex_width * tex_height * 4, nullptr, GL_DYNAMIC_STORAGE_BIT);
+        glBindBuffer(GL_PIXEL_PACK_BUFFER, 0);
+        CUDA_API_CHECK(cudaGraphicsGLRegisterBuffer(&pbo_res, pbo, cudaGraphicsRegisterFlagsReadOnly));
+
+        return true;
+    }
+
+    void download_pixels() {
+        glBindFramebuffer(GL_READ_FRAMEBUFFER, fbo);
+        glBindBuffer(GL_PIXEL_PACK_BUFFER, pbo);
+        glReadPixels(0, 0, tex_width, tex_height, GL_BGRA, GL_UNSIGNED_INT_8_8_8_8_REV, (void *) 0);
+        glBindBuffer(GL_PIXEL_PACK_BUFFER, 0);
+    }
+
+private:
+
+    static bool check_frame_buffer() {
+        auto status = glCheckFramebufferStatus(GL_FRAMEBUFFER);
+        if (status != GL_FRAMEBUFFER_COMPLETE) [[unlikely]] {
+            SPDLOG_ERROR("Framebuffer is not complete 0x{:x}.", status);
+            RET_ERROR;
+        }
+        return true;
+    }
+
+
+};
+
+#endif //REMOTEAR2_FRAME_BUFFER_HELPER_HPP

+ 86 - 9
src/main.cpp

@@ -1,7 +1,9 @@
 #include "augment_renderer.h"
 #include "config.h"
+#include "frame_buffer_helper.hpp"
 #include "stereo_camera.hpp"
 #include "texture_renderer.h"
+#include "video_encoder.h"
 
 #include <imgui.h>
 #include <imgui_impl_glfw.h>
@@ -10,12 +12,19 @@
 #include <glad/gl.h>
 #include <GLFW/glfw3.h>
 
+#include <fmt/chrono.h>
+#include <fmt/format.h>
 #include <spdlog/spdlog.h>
 
 #include <cassert>
+#include <cstdlib>
+
+CUcontext cuda_ctx;
 
 int main() {
 
+    spdlog::set_level(spdlog::level::trace);
+
     // setup glfw and main window
     glfwSetErrorCallback([](int error, const char *desc) {
         SPDLOG_ERROR("GLFW error: code = {}, description = {}", error, desc);
@@ -57,16 +66,32 @@ int main() {
     ImGui_ImplGlfw_InitForOpenGL(main_window, true);
     ImGui_ImplOpenGL3_Init();
 
+    // setup cuda context
+    cuInit(0);
+    create_cuda_context(&cuda_ctx);
+
     // working staffs
     stereo_camera camera;
     texture_renderer tex_renderer;
 
+    frame_buffer_helper fbo_helper;
+    fbo_helper.initialize(image_width * 2, image_height);
+
+    video_encoder encoder;
+    encoder.initialize();
+
     augment_renderer left_ar, right_ar;
     left_ar.initialize(&tex_renderer);
     right_ar.initialize(&tex_renderer);
     left_ar.set_background(&camera.left_rgb_image);
     right_ar.set_background(&camera.right_rgb_image);
 
+    int camera_fps = default_camera_fps;
+    float exposure_time_ms = default_camera_exposure_time_ms;
+    float analog_gain = default_camera_analog_gain;
+
+    FILE *video_save_file = nullptr;
+
     // main loop
     while (!glfwWindowShouldClose(main_window)) {
 
@@ -80,12 +105,14 @@ int main() {
 
         if (ImGui::Begin("Remote AR Control")) {
 
+            // extra actions to make consistency
+            if (!camera.is_capturing() && encoder.is_encoding()) {
+                encoder.stop_encode();
+            }
+
             // camera control
             if (ImGui::CollapsingHeader("Camera")) {
-
-                static int trigger_interval_ms = default_camera_capture_interval_ms;
-                static float exposure_time_ms = default_camera_exposure_time_ms;
-                static float analog_gain = default_camera_analog_gain;
+                ImGui::PushID("Camera");
 
                 // camera actions
                 ImGui::SeparatorText("Actions");
@@ -100,7 +127,7 @@ int main() {
                     ImGui::SameLine();
                     if (!camera.is_capturing()) {
                         if (ImGui::Button("Start")) {
-                            camera.start_capture(1000 * exposure_time_ms, analog_gain, trigger_interval_ms);
+                            camera.start_capture(1000 * exposure_time_ms, analog_gain, camera_fps);
                         }
                     } else {
                         if (ImGui::Button("Stop")) {
@@ -119,9 +146,9 @@ int main() {
                     }
 
                     ImGui::PushItemWidth(200);
-                    ImGui::SliderInt("Trigger Interval (ms)", &trigger_interval_ms, 17, 100);
+                    ImGui::SliderInt("Frame Rate (fps)", &camera_fps, 1, 60);
                     ImGui::DragFloat("Exposure Time (ms)", &exposure_time_ms,
-                                     0.1, 1, (float) trigger_interval_ms, "%.1f");
+                                     0.1, 1, 1e3f / (float) camera_fps, "%.1f");
                     ImGui::DragFloat("Analog Gain (dB)", &analog_gain, 0.1, 0, 24, "%.1f");
                     ImGui::PopItemWidth();
 
@@ -129,6 +156,34 @@ int main() {
                         ImGui::EndDisabled();
                     }
                 }
+
+                ImGui::PopID();
+            }
+
+            // video streamer control
+            if (camera.is_capturing() && ImGui::CollapsingHeader("Video Streamer")) {
+                ImGui::PushID("Streamer");
+
+                ImGui::SeparatorText("Actions");
+                if (!encoder.is_encoding()) {
+                    if (ImGui::Button("Start")) {
+                        // create save file
+                        auto file_name = fmt::format("record_{:%Y_%m_%d_%H_%M_%S}.hevc",
+                                                     std::chrono::system_clock::now());
+                        video_save_file = fopen(file_name.c_str(), "wb");
+
+                        encoder.start_encode(fbo_helper.tex_width, fbo_helper.tex_height, camera_fps);
+                        SPDLOG_INFO("Video streamer started.");
+                    }
+                } else {
+                    if (ImGui::Button("Close")) {
+                        encoder.stop_encode();
+                        fclose(video_save_file);
+                        SPDLOG_INFO("Video streamer stopped.");
+                    }
+                }
+
+                ImGui::PopID();
             }
 
         }
@@ -136,6 +191,7 @@ int main() {
         ImGui::Render();
 
         int frame_width, frame_height;
+        glBindFramebuffer(GL_DRAW_FRAMEBUFFER, 0);
         glfwGetFramebufferSize(main_window, &frame_width, &frame_height);
         glViewport(0, 0, frame_width, frame_height);
         glClear(GL_COLOR_BUFFER_BIT);
@@ -144,19 +200,40 @@ int main() {
             camera.retrieve_raw_images();
             camera.debayer_images();
 
-            augment_renderer::render_config config{-1, 1, 2, -2};
-            left_ar.render(&config);
+            // draw frame in the screen
+            left_ar.render({-1, 1, 2, -2});
         }
 
         ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());
         glfwSwapBuffers(main_window);
 
+        if (encoder.is_encoding()) {
+            // draw frame for streaming
+            glBindFramebuffer(GL_DRAW_FRAMEBUFFER, fbo_helper.fbo);
+            glViewport(0, 0, fbo_helper.tex_width, fbo_helper.tex_height);
+            left_ar.render({-1, -1, 1, 2});
+            right_ar.render({0, -1, 1, 2});
+
+            // encode frame
+            fbo_helper.download_pixels();
+            void *frame_data;
+            size_t frame_length;
+            encoder.encode_frame(fbo_helper.pbo_res, &frame_data, &frame_length);
+
+            // save encoded frame
+            fwrite(frame_data, frame_length, 1, video_save_file);
+        }
+
         if (camera.is_capturing()) {
             glFlush();
         }
     }
 
     // cleanup
+    std::atexit([]() {
+        cuCtxDestroy(cuda_ctx);
+    });
+
     ImGui_ImplOpenGL3_Shutdown();
     ImGui_ImplGlfw_Shutdown();
     ImGui::DestroyContext();

+ 6 - 10
src/stereo_camera.hpp

@@ -73,14 +73,10 @@ struct stereo_camera {
         assert(is_capturing());
 
         // clean old images
-        if (left_raw_image != nullptr) {
-            delete left_raw_image;
-            left_raw_image = nullptr;
-        }
-        if (right_raw_image != nullptr) {
-            delete right_raw_image;
-            right_raw_image = nullptr;
-        }
+        delete left_raw_image;
+        delete right_raw_image;
+        left_raw_image = nullptr;
+        right_raw_image = nullptr;
 
         // retrieve new images
         left_camera.retrieve_image(&left_raw_image);
@@ -125,10 +121,10 @@ private:
         return true;
     }
 
-    void start_trigger_thread(int trigger_interval_ms) {
+    void start_trigger_thread(int fps) {
         assert(trigger_thread == nullptr);
         trigger_thread = new std::thread{[=, this]() {
-            auto trigger_interval = std::chrono::milliseconds{trigger_interval_ms};
+            auto trigger_interval = std::chrono::microseconds{(int) 1e6 / fps};
             auto next_trigger_time = std::chrono::high_resolution_clock::now();
             while (true) {
                 if (should_stop.test()) break;

+ 226 - 0
src/video_encoder.cpp

@@ -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, &reg_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;
+}

+ 33 - 0
src/video_encoder.h

@@ -0,0 +1,33 @@
+#ifndef REMOTEAR2_VIDEO_ENCODER_H
+#define REMOTEAR2_VIDEO_ENCODER_H
+
+#include <cuda_gl_interop.h>
+
+#include <memory>
+
+class video_encoder {
+public:
+
+    video_encoder();
+
+    ~video_encoder();
+
+    bool initialize();
+
+    bool start_encode(int width, int height, int fps);
+
+    void stop_encode();
+
+    bool encode_frame(void *frame_ptr, void **output_ptr, size_t *output_size);
+
+    bool encode_frame(cudaGraphicsResource *res, void **output_ptr, size_t *output_size);
+
+    bool is_encoding();
+
+private:
+    struct impl;
+    std::unique_ptr<impl> pimpl;
+};
+
+
+#endif //REMOTEAR2_VIDEO_ENCODER_H