DepthGuide/src/impl/main_impl.cpp

#include "main_impl.h"
#include "device/orb_camera_ui.h"
#include "core/image_utility.hpp"
#include "module/image_streamer.h"
#include "module/image_viewer.h"
#include "module/viewport_downloader.hpp"
#include "object_names.h"

#include <boost/asio/io_context.hpp>
#include <boost/asio/post.hpp>
#include <boost/asio/steady_timer.hpp>

#include <glad/gl.h>
#include <GLFW/glfw3.h>

#include <imgui.h>
#include <imgui_impl_glfw.h>
#include <imgui_impl_opengl3.h>

// make glad happy
#include "core/imgui_utility.hpp"

using boost::asio::io_context;
using boost::asio::post;
using boost::asio::steady_timer;
using boost::system::error_code;

CUcontext cuda_ctx = nullptr;
GLFWwindow *window = nullptr;
smart_cuda_stream *default_cuda_stream = nullptr;
io_context *main_ctx;
object_manager *main_ob;

std::unique_ptr<steady_timer> ui_timer;
std::chrono::milliseconds ui_interval;

// modules
std::unique_ptr<orb_camera_ui> orb_cam;
std::unique_ptr<image_viewer> bg_viewer; // background viewer
std::unique_ptr<viewport_downloader> out_downloader;
std::unique_ptr<image_streamer> out_streamer; // output streamer

void init_cuda() {
    cuInit(0);

    auto cuda_dev = CUdevice();
    CUDA_API_CHECK(cuDeviceGet(&cuda_dev, 0)); // TODO: select device
    CUDA_API_CHECK(cuCtxCreate(&cuda_ctx, CU_CTX_SCHED_AUTO, cuda_dev));

    default_cuda_stream = new smart_cuda_stream();

    std::atexit([] { // elegant cleanup
        cuCtxDestroy(cuda_ctx);
    });
}

void init_window() {
    // set GLFW error handler
    glfwSetErrorCallback([](int error, const char *desc) {
        SPDLOG_ERROR("GLFW error: code = {}, description = {}", error, desc);
    });

    // create main window
    auto ret = glfwInit();
    assert(ret == GLFW_TRUE);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 6);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
    window = glfwCreateWindow(800, 600, "DepthGuide V1.-1", nullptr, nullptr); // TODO: select width and height
    assert(window != nullptr);
    glfwMakeContextCurrent(window);
    glfwSwapInterval(0);

    // load opengl functions
    auto version = gladLoadGL(glfwGetProcAddress);
    assert(version > 0);
    SPDLOG_INFO("Loaded OpenGL {}.{}", GLAD_VERSION_MAJOR(version), GLAD_VERSION_MINOR(version));

    // enable color blending
    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);

#ifndef NDEBUG
    // log opengl error
    glEnable(GL_DEBUG_OUTPUT);
    glDebugMessageCallback([](GLenum source, GLenum type, GLuint id, GLenum severity,
                              GLsizei length, const GLchar *message, const void *user_data) {
        if (type == GL_DEBUG_TYPE_ERROR) {
            SPDLOG_ERROR("OpenGL error: type = {}, severity = {}, message = {}", type, severity, message);
            assert(false);
        }
    }, nullptr);
#endif

    // setup imgui context
    IMGUI_CHECKVERSION();
    ImGui::CreateContext();
    auto io = ImGui::GetIO();
    io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;
    ImGui::StyleColorsDark();
    ImGui_ImplGlfw_InitForOpenGL(window, true);
    ImGui_ImplOpenGL3_Init();

    // elegant cleanup
    std::atexit([] {
        ImGui_ImplOpenGL3_Shutdown();
        ImGui_ImplGlfw_Shutdown();
        ImGui::DestroyContext();

        glfwDestroyWindow(window);
        glfwTerminate();
    });
}

void init_om() {
    main_ctx = new io_context();
    main_ob = new object_manager({.ctx = main_ctx});
    OBJ_SAVE(img_color, image_u8c3());
    OBJ_SAVE(img_depth, image_f32c1());
    OBJ_SAVE(img_bg, image_u8c3());
    OBJ_SAVE(img_out, image_u8c4());

    OBJ_SIG(img_color)->connect(INT_MIN, [=](obj_name_type _) {
        OBJ_SAVE(img_bg, OBJ_QUERY(image_u8c3, img_color));
    });
}

void init_modules() {
    auto orb_cam_conf = orb_camera_ui::create_config{
            .cf_name = img_color, .df_name = img_depth,
            .stream = default_cuda_stream,
    };
    orb_cam = std::make_unique<orb_camera_ui>(orb_cam_conf);

    auto bg_viewer_conf = image_viewer::create_config{
            .mode = VIEW_COLOR_DEPTH, .flip_y = true,
            .stream = default_cuda_stream,
    };
    auto &bg_extra_conf = bg_viewer_conf.extra.color_depth;
    bg_extra_conf.c_name = img_color;
    bg_extra_conf.d_name = img_depth;
    bg_viewer = std::make_unique<image_viewer>(bg_viewer_conf);

    auto out_down_conf = viewport_downloader::create_config{
            .type = PIX_RGBA, .stream = default_cuda_stream
    };
    out_downloader = std::make_unique<viewport_downloader>(out_down_conf);

    auto out_streamer_conf = image_streamer::create_config{
            .img_name = img_out, .cuda_ctx = &cuda_ctx, .stream = default_cuda_stream
    };
    out_streamer = std::make_unique<image_streamer>(out_streamer_conf);
}

void ui_timer_func(error_code ec) {
    if (ec == boost::asio::error::operation_aborted) return;
    assert(ec == error_code());
    show_ui();
    ui_timer->expires_after(ui_interval);
    ui_timer->async_wait(ui_timer_func);
}

void init_all() {
    init_cuda();
    init_window();
    init_om();
    init_modules();

    ui_interval = std::chrono::milliseconds(33); // TODO: select refresh rate
    ui_timer = std::make_unique<steady_timer>(*main_ctx, ui_interval);
    ui_timer->async_wait(ui_timer_func);
}

void show_ui() {
    glfwPollEvents();
    ImGui_ImplOpenGL3_NewFrame();
    ImGui_ImplGlfw_NewFrame();
    ImGui::NewFrame();

    if (glfwWindowShouldClose(window)) {
        ui_timer->cancel();
        main_ctx->stop();
        return;
    }

    if (ImGui::Begin("Depth Guide Control")) {
        ImGui::PushItemWidth(200);

        if (ImGui::CollapsingHeader("Camera")) {
            auto id_guard = imgui_id_guard("camera");
            orb_cam->show();
        }

        if (ImGui::CollapsingHeader("Streamer")) {
            auto id_guard = imgui_id_guard("streamer");
            out_streamer->show();
        }

        if (ImGui::CollapsingHeader("Debug")) {
            if (ImGui::TreeNode("Background")) {
                bg_viewer->show();
                ImGui::TreePop();
            }
            if (ImGui::TreeNode("Memory Pool")) {
                if (ImGui::Button("Purge")) {
                    post(*main_ctx, [] { global_mp.purge(); });
                }
                ImGui::TreePop();
            }
        }

        ImGui::PopItemWidth();
    }
    ImGui::End();
    ImGui::Render();

    cv::Size frame_size;
    glBindFramebuffer(GL_DRAW_FRAMEBUFFER, 0);
    glfwGetFramebufferSize(window, &frame_size.width, &frame_size.height);
    glViewport(0, 0, frame_size.width, frame_size.height);
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

    bg_viewer->render();

    // TODO: for test
    auto bg_img = out_downloader->download_rgba();
    OBJ_SAVE(img_out, bg_img);

    ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());
    glfwSwapBuffers(window);
}

void cleanup() {
    ui_timer = nullptr;
    orb_cam = nullptr;
    bg_viewer = nullptr;
    out_downloader = nullptr;
    out_streamer = nullptr;

    delete main_ob;
    delete main_ctx;
}