RemoteAR3/src/main_ext.cpp

#include "components/registration.h"
#include "core/local_connection.h"
#include "cuda_helper.hpp"
#include "experiment/probe_p2p.h"
#include "imgui_utility.h"
#include "main_impl/impl_types.h"
#include "simple_mq.h"
#include "utility.hpp"
#include "variable_defs.h"
#include "vis_marker_detector.h"
#include "vtk_viewer.h"

#ifdef _MSC_VER

#include <format>

#define fmt std

#endif

#include <spdlog/spdlog.h>

#include <yaml-cpp/yaml.h>

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

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

#include <cxxopts.hpp>

#include <boost/iostreams/device/mapped_file.hpp>

#include <cassert>
#include <fstream>
#include <thread>
#include <queue>
#include <vector>

using namespace simple_mq_singleton;
using namespace sophiar;

using boost::iostreams::mapped_file;

std::ofstream main_log_file;
log_timer global_timer;

// global variable definition
CUcontext cuda_ctx = nullptr;
int main_window_width = -1, main_window_height = -1;
GLFWwindow *main_window = nullptr;

float process_frame_rate = 0;

// camera related
extern bool mono_mode;

// render related
extern bool enable_augment_without_nav;
extern std::unique_ptr<scene_render> augment_render;
extern std::vector<augment_store_type> augment_items;

std::string sophiar_config_path;
std::unique_ptr<std::thread> sophiar_thread;
local_connection sophiar_conn;
bool is_tracking = false;

bool enable_reg = false;
std::string probe_model_path;
std::vector<registration_target> reg_targets;
std::unique_ptr<registration> reg;

std::unique_ptr<probe_p2p> exp_probe;

bool debug_options = false;
bool show_vtk_debug = false;
bool show_imgui_demo = false;
std::unique_ptr<vtk_viewer> vtk_test1, vtk_test2;

std::queue<std::function<void()>> simple_eq;
std::queue<std::function<void()>> close_funcs;

std::unique_ptr<camera_related> left;
std::unique_ptr<camera_related> right;

void initialize_render();

void initialize_augment_accuracy();

void initialize_main_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);
    main_window = glfwCreateWindow(main_window_width, main_window_height,
                                   "RemoteAR V3.-1", nullptr, nullptr);
    assert(main_window != nullptr);
    glfwMakeContextCurrent(main_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(main_window, true);
    ImGui_ImplOpenGL3_Init();

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

        glfwDestroyWindow(main_window);
        glfwTerminate();
    });
}

void load_camera_config(YAML::Node conf);

void load_encoder_config(YAML::Node conf);

void load_sender_config(YAML::Node conf);

void load_config(const std::string &config_path) {
    auto conf = YAML::LoadFile(config_path);
    SPDLOG_INFO("Use config file: {}", config_path);

    // initialize mq
    mq();

    // load module configs
    load_camera_config(conf["camera"]);
    load_encoder_config(conf["output"]);
    load_sender_config(conf["sender"]);

    // load render configs
    auto opengl_conf = conf["opengl"];
    auto shader_dir = opengl_conf["shader_dir"].as<std::string>();
    mq().update_variable(SHADER_DIR, shader_dir);

    // load main window config
    auto window_conf = conf["main_window"];
    main_window_width = window_conf["width"].as<int>();
    main_window_height = window_conf["height"].as<int>();

    // load sophiar config
    auto sophiar_conf = conf["sophiar"];
    sophiar_config_path = sophiar_conf["config"].as<std::string>();
    left->trans_var = sophiar_conf["left_camera_trans_var"].as<std::string>();
    right->trans_var = sophiar_conf["right_camera_trans_var"].as<std::string>();
    probe_model_path = sophiar_conf["probe_model"].as<std::string>();

    // load augment items
    for (auto item: sophiar_conf["augment"]) {
        auto &store = augment_items.emplace_back();
        store.name = item["name"].as<std::string>();
        store.trans_var = item["trans_var"].as<std::string>();
        store.model_path = item["model_file"].as<std::string>();
        if (auto bg_conf = item["background"]; bg_conf) {
            store.background = bg_conf.as<std::string>();
        } else {
            store.background = {};
        }
        if (auto reg_conf = item["registration"]; reg_conf) {
            reg_targets.emplace_back(
                    store.name, store.model_path,
                    std::vector<Eigen::Vector3d>{},
                    reg_conf["target_var"].as<std::string>(),
                    reg_conf["collect_var"].as<std::string>(),
                    reg_conf["collect_obj"].as<std::string>(),
                    reg_conf["probe_var"].as<std::string>());
        }
    }

    // make variables exist
    mq().update_variable(SENDER_CONNECTED, false);
    mq().update_variable(REQUEST_IDR, false);
    mq().update_variable(ENCODER_SHOULD_RESET, false);
}

void initialize_main(int argc, char *argv[]) {

    // parse arguments
    auto options = cxxopts::Options{"RemoteAR3", "Capture surgery scene and do augmentation works."};
    options.add_options()
            ("c,config", "Config file", cxxopts::value<std::string>()->default_value("config.yaml"));
    auto opts = options.parse(argc, argv);

    // initialize cuda
    constexpr auto default_cuda_device_id = 0;
    cuInit(0);
    int cuda_device_count;
    CUDA_API_CHECK(cuDeviceGetCount(&cuda_device_count));
    assert(cuda_device_count > default_cuda_device_id);
    CUdevice cuda_device;
    CUDA_API_CHECK(cuDeviceGet(&cuda_device, default_cuda_device_id));
    CUDA_API_CHECK(cuCtxCreate(&cuda_ctx, CU_CTX_SCHED_AUTO, cuda_device));
    mq().update_variable(CUDA_CONTEXT, cuda_ctx);
    std::atexit([] { // elegant cleanup
        cuCtxDestroy(cuda_ctx);
    });

    left = std::make_unique<camera_related>();
    right = std::make_unique<camera_related>();

    load_config(opts["config"].as<std::string>());
    initialize_main_window();

    // initialize modules
    initialize_render();
    initialize_augment_accuracy();

    // initialize sophiar
    assert(sophiar_thread == nullptr);
    sophiar_thread = std::make_unique<std::thread>([=] {
        run_sophiar(sophiar_config_path);
    });

    // initialize components
    reg.reset(registration::create({&sophiar_conn,
                                    probe_model_path}));
    for (auto &item: reg_targets) {
        reg->add_target(item);
    }
    exp_probe.reset(probe_p2p::create(
            {"./models/Probe.stl",
             "probe_in_tracker"}));

    // initialize vtk test viewer
    vtk_test1 = std::make_unique<vtk_viewer>();
    vtk_test2 = std::make_unique<vtk_viewer>();
//    vtk_test1->clear_actor(); // TODO: remove them
//    vtk_test1->add_actor(create_actor("./models/femur.stl"));
//    vtk_test1->reset_camera();
//    vtk_test2->add_actor(create_actor("./models/tibia.stl"));
//    vtk_test1->start_picking();
}

void start_tracking() {
    CALL_CHECK(sophiar_conn.start_object("tracker_all"));
    is_tracking = true;
}

void close_cameras();

void stop_encoder();

void stop_sender();

void cleanup() {
    close_cameras();
    stop_encoder();
    stop_sender();

    // avoid cudaErrorCudartUnloading
    left.reset();
    right.reset();

    // stop sophiar
    assert(sophiar_thread != nullptr);
    stop_sophiar();
    sophiar_thread->join();
    sophiar_thread.reset();

    // cleanup modules
    while (!close_funcs.empty()) {
        close_funcs.front()();
        close_funcs.pop();
    }
}

void show_camera_ui();

void show_encoder_ui();

void show_sender_ui();

void show_augment_ui();

void show_augment_accuracy_ui();

void prepare_imgui_frame() {

    glfwPollEvents();
    ImGui_ImplOpenGL3_NewFrame();
    ImGui_ImplGlfw_NewFrame();
    ImGui::NewFrame();

    if (show_imgui_demo) {
        ImGui::ShowDemoWindow();
    }

    if (ImGui::Begin("Remote AR Control")) {
        ImGui::PushItemWidth(200);

        if (ImGui::CollapsingHeader("Camera")) {
            ImGui::PushID("Camera");
            show_camera_ui();
            ImGui::PopID();
        }

        if (ImGui::CollapsingHeader("Navigation")) {
            ImGui::PushID("Encoder");

            ImGui::SeparatorText("Actions");
            if (ImGui::Button("Start Tracking")) {
                simple_eq.emplace(start_tracking);
            }
            if (is_tracking) {
                ImGui::Checkbox("Registration Panel", &enable_reg);
            }

            ImGui::SeparatorText("Infos");
            auto helper_func = [&](const std::string &var_name, const std::string &show_name, bool last = false) {
                auto var = sophiar_conn.query_transform_variable(var_name);
                bool var_ok = var.has_value();
                if (var_ok) {
                    ImGui::PushStyleColor(ImGuiCol_Text, (ImVec4) ImColor(0, 255, 0));
                } else {
                    ImGui::PushStyleColor(ImGuiCol_Text, (ImVec4) ImColor(255, 0, 0));
                }
                ImGui::Checkbox(show_name.c_str(), &var_ok);
                ImGui::PopStyleColor();
                if (!last) {
                    ImGui::SameLine();
                }
            };
//            ImGui::BeginDisabled();
            helper_func("camera_ref_in_tracker", "Camera"); // TODO set in config
            helper_func("probe_in_tracker", "Probe");
            helper_func("femur_ref_in_tracker", "Femur");
            helper_func("tibia_ref_in_tracker", "Tibia", true);
//            ImGui::EndDisabled();

            ImGui::PopID();
        }

        if (ImGui::CollapsingHeader("Augment Render")) {
            ImGui::PushID("Augment");
            show_augment_ui();
            ImGui::PopID();
        }

        if (ImGui::CollapsingHeader("Video Encoder")) {
            ImGui::PushID("Encoder");
            show_encoder_ui();
            ImGui::PopID();
        }

        if (ImGui::CollapsingHeader("Frame Sender")) {
            ImGui::PushID("Sender");
            show_sender_ui();
            ImGui::PopID();
        }

        if (ImGui::CollapsingHeader("Debug")) {
            ImGui::PushID("Debug");
            ImGui::Checkbox("Debug VTK Viewer", &show_vtk_debug);
            ImGui::Checkbox("Show ImGui Demo", &show_imgui_demo);
            ImGui::Checkbox("Debug Options", &debug_options);
            ImGui::Checkbox("Enable Augment without Navigation",
                            &enable_augment_without_nav);
            ImGui::PopID();
        }

        if (ImGui::CollapsingHeader("Experiment")) {
            ImGui::PushID("Experiment");
            if (mono_mode) {
                if (ImGui::TreeNode("Augment Accuracy")) {
                    show_augment_accuracy_ui();
                    ImGui::TreePop();
                }
            }
            if (ImGui::TreeNode("Probe P2P")) {
                exp_probe->show_ui();
                ImGui::TreePop();
            }
            ImGui::PopID();
        }

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

    if (show_vtk_debug) {
        if (ImGui::Begin("VTK Test A", nullptr, vtk_viewer::no_scroll_flag)) {
            vtk_test1->show();
        }
        ImGui::End();
        if (ImGui::Begin("VTK Test B", nullptr, vtk_viewer::no_scroll_flag)) {
            vtk_test2->show();
        }
        ImGui::End();
    }

    if (enable_reg) {
        reg->show();
    }

    ImGui::Render();
}

void process_camera_frames();

void prepare_augment_info();

void process_augment_accuracy();

void process_frame() {
    prepare_augment_info();
    process_augment_accuracy();
    exp_probe->process();
    process_camera_frames();
}

void handle_imgui_events() {
    while (!simple_eq.empty()) {
        simple_eq.front()();
        simple_eq.pop();
    }

    // registration related
    if (enable_reg) {
        reg->process();
    }
}