Implemented registration.

CMakeLists.txt

@@ -8,7 +8,10 @@ add_executable(RemoteAR3 src/main.cpp
         src/frame_sender.cpp
         src/image_process.cpp
         src/simple_mq.cpp
-        src/third_party/rs.c)
+        src/third_party/rs.c
+        src/components/registration.cpp)
+
+target_include_directories(${PROJECT_NAME} PRIVATE ./src)
 
 add_subdirectory(src/image_process)
 target_link_libraries(${PROJECT_NAME} ImageProcess)
@@ -108,7 +111,6 @@ target_include_directories(${PROJECT_NAME} PRIVATE ${Boost_INCLUDE_DIRS})
 target_link_libraries(${PROJECT_NAME} ${Boost_LIBRARIES})
 
 # VTK config
-set(VTK_DIR /home/tpx/src/VTK-9.3.0/Build/lib/cmake/vtk-9.3)
 find_package(VTK REQUIRED)
 target_link_libraries(${PROJECT_NAME} ${VTK_LIBRARIES})
 vtk_module_autoinit(TARGETS ${PROJECT_NAME} MODULES ${VTK_LIBRARIES})

data/config.yaml

@@ -6,9 +6,13 @@ camera:
     width: 2491
     height: 2077
     angle: 32.351172
+    #    width: 2458
+    #    height: 2048
+    #    angle: 32.36803962255573
     data:
       left: ./left_proj.dat
       right: ./right_proj.dat
+  #      right: /home/tpx/project/RemoteAR2/data/right_proj.dat
   capture:
     frame_rate: 40
     expo_time_ms: 12
@@ -32,10 +36,21 @@ sophiar:
   config: ./sophiar_config.json
   left_camera_trans_var: left_camera_in_tracker
   right_camera_trans_var: right_camera_in_tracker
+  probe_model: /home/tpx/data/stls/Probe.stl
   augment:
     - name: Femur
       trans_var: femur_in_tracker
       stl_file: /home/tpx/project/RemoteAR3/data/femur.stl
+      registration:
+        collect_obj: point_picker_in_femur_ref
+        collect_var: picked_point_in_femur_ref
+        target_var: femur_in_femur_ref
+        probe_var: probe_in_femur
     - name: Tibia
       trans_var: tibia_in_tracker
-      stl_file: /home/tpx/project/RemoteAR3/data/tibia.stl
+      stl_file: /home/tpx/project/RemoteAR3/data/tibia.stl
+      registration:
+        collect_obj: point_picker_in_tibia_ref
+        collect_var: picked_point_in_tibia_ref
+        target_var: tibia_in_tibia_ref
+        probe_var: probe_in_tibia

data/sophiar_config.json

@@ -11,7 +11,7 @@
     {
       "name": "femur_in_femur_ref",
       "type": "transform_obj",
-      "value": [
+      "value_": [
         49.7073,
         -72.6173,
         -9.09289,
@@ -28,7 +28,7 @@
     {
       "name": "tibia_in_tibia_ref",
       "type": "transform_obj",
-      "value": [
+      "value_": [
         -3.91566,
         -85.3472,
         -42.0667,
@@ -244,7 +244,7 @@
       "type": "ndi_interface",
       "name": "ndi",
       "init_config": {
-        "address_type": "ethernet",
+        "address_type": "serial",
         "ip": "10.0.0.5",
         "tcp_port": 8765,
         "com_port": "/dev/ttyUSB0",

src/components/registration.cpp

@@ -0,0 +1,457 @@
+#include "registration.h"
+#include "vtk_viewer.h"
+#include "utility.hpp"
+
+#include <vtkActor.h>
+#include <vtkCellLocator.h>
+#include <vtkIterativeClosestPointTransform.h>
+#include <vtkLandmarkTransform.h>
+#include <vtkMatrix4x4.h>
+#include <vtkNamedColors.h>
+#include <vtkPolyData.h>
+#include <vtkProperty.h>
+#include <vtkSmartPointer.h>
+
+#include <queue>
+#include <vector>
+
+struct registration::impl {
+
+    static constexpr auto MIN_REG_POINTS = 3;
+    static constexpr auto DIS_LIMIT_1 = 0.5;
+    static constexpr auto DIS_LIMIT_2 = 1.0;
+    static constexpr auto NAN_VALUE = std::numeric_limits<float>::quiet_NaN();
+
+    enum show_mode {
+        CONFIG,
+        COLLECTING,
+        FINISHED
+    };
+
+    struct target_store_type {
+        std::string name;
+        vtkSmartPointer<vtkPolyData> model;
+        vtkSmartPointer<vtkActor> model_actor;
+        vtkSmartPointer<vtkCellLocator> model_locator;
+        std::unique_ptr<smart_point_sets> all_points;
+        std::unique_ptr<smart_point_sets> pending_points;
+        std::unique_ptr<smart_point_sets> current_point;
+        std::unique_ptr<smart_point_sets> finished_points;
+        std::unique_ptr<smart_point_sets> level_points[3];
+
+        std::string target_var_name;
+        std::string point_var_name;
+        std::string collect_obj_name;
+        std::string probe_var_name;
+
+        bool is_finished = false;
+        float max_error = 0;
+    };
+
+    int cur_target_id = -1;
+    std::vector<target_store_type> targets;
+    target_store_type *cur_target = nullptr;
+
+    std::unique_ptr<vtk_viewer> viewer;
+    bool is_picking = false;
+    bool is_collecting = false;
+    std::vector<Eigen::Vector3d> collected_points;
+    std::queue<std::function<void()>> eq;
+
+    sophiar::local_connection *conn = nullptr;
+    vtkSmartPointer<vtkActor> probe_actor;
+    float probe_test_error = NAN_VALUE;
+
+    // colors
+    vtkColor4d target_color;
+    vtkColor4d idle_color;
+    vtkColor4d pending_color;
+    vtkColor4d current_color;
+    vtkColor4d finished_color;
+    vtkColor4d level_color[3]; // - 0.5mm - 1.0mm -
+
+    impl() {
+        viewer = std::make_unique<vtk_viewer>();
+
+        // preload colors
+        vtkNew<vtkNamedColors> colors;
+        target_color = colors->GetColor4d("silver");
+        idle_color = colors->GetColor4d("white");
+        pending_color = colors->GetColor4d("red");
+        current_color = colors->GetColor4d("yellow");
+        finished_color = colors->GetColor4d("lime");
+        level_color[0] = colors->GetColor4d("lime");
+        level_color[1] = colors->GetColor4d("yellow");
+        level_color[2] = colors->GetColor4d("red");
+    }
+
+    void switch_viewer_mode(show_mode mode) {
+        assert(cur_target != nullptr);
+        viewer->clear_actor();
+        viewer->add_actor(cur_target->model_actor);
+        viewer->add_actor(probe_actor);
+        switch (mode) {
+            case CONFIG: {
+                viewer->add_actor(cur_target->all_points->get_actor());
+                break;
+            }
+            case COLLECTING: {
+                viewer->add_actor(cur_target->pending_points->get_actor());
+                viewer->add_actor(cur_target->current_point->get_actor());
+                viewer->add_actor(cur_target->finished_points->get_actor());
+                break;
+            }
+            case FINISHED: {
+                for (auto &points: cur_target->level_points) {
+                    viewer->add_actor(points->get_actor());
+                }
+                break;
+            }
+        }
+    }
+
+    void add_target(const registration_target &conf) {
+        auto &target = targets.emplace_back();
+        target.name = conf.name;
+        target.model = load_stl(conf.model_path);
+        target.model_actor = create_actor(target.model);
+        target.model_locator = vtkSmartPointer<vtkCellLocator>::New();
+        target.model_locator->SetDataSet(target.model);
+        target.model_locator->BuildLocator();
+
+        // copy information
+        target.target_var_name = conf.target_var_name;
+        target.collect_obj_name = conf.collect_obj_name;
+        target.point_var_name = conf.collect_var_name;
+        target.probe_var_name = conf.probe_var_name;
+
+        // create point sets
+        target.all_points = std::make_unique<smart_point_sets>();
+        target.pending_points = std::make_unique<smart_point_sets>();
+        target.current_point = std::make_unique<smart_point_sets>();
+        target.finished_points = std::make_unique<smart_point_sets>();
+        for (auto &points: target.level_points) {
+            points = std::make_unique<smart_point_sets>();
+        }
+
+        // set colors
+        target.model_actor->GetProperty()->SetColor(target_color.GetData());
+        target.all_points->get_actor()->GetProperty()->SetColor(idle_color.GetData());
+        target.pending_points->get_actor()->GetProperty()->SetColor(pending_color.GetData());
+        target.current_point->get_actor()->GetProperty()->SetColor(current_color.GetData());
+        target.finished_points->get_actor()->GetProperty()->SetColor(finished_color.GetData());
+        for (auto k = 0; k < 3; ++k) {
+            target.level_points[k]->get_actor()->GetProperty()->SetColor(level_color[k].GetData());
+        }
+    }
+
+    void change_target() {
+        assert(cur_target_id != -1);
+        cur_target = &targets[cur_target_id];
+        if (cur_target->is_finished) {
+            switch_viewer_mode(FINISHED);
+        } else {
+            switch_viewer_mode(CONFIG);
+        }
+        viewer->reset_camera();
+    }
+
+    bool progress_reg_point() {
+        if (!cur_target->current_point->empty()) {
+            cur_target->finished_points->add_point(
+                    cur_target->current_point->pop_front());
+        }
+        if (cur_target->pending_points->empty()) return false;
+        cur_target->current_point->add_point(
+                cur_target->pending_points->pop_front());
+        return true;
+    }
+
+    void start() {
+        assert(cur_target != nullptr);
+        assert(!is_collecting);
+        is_collecting = true;
+        cur_target->is_finished = false;
+        conn->mark_variable_disposal(cur_target->point_var_name);
+        CALL_CHECK(conn->start_object(cur_target->collect_obj_name));
+        collected_points.clear();
+
+        // copy points
+        cur_target->all_points->for_each([this](void *, const Eigen::Vector3d &point) {
+            cur_target->pending_points->add_point(point);
+        });
+        progress_reg_point();
+
+        // switch actor
+        switch_viewer_mode(COLLECTING);
+    }
+
+    void stop() {
+        assert(is_collecting);
+        is_collecting = false;
+        CALL_CHECK(conn->stop_object(cur_target->collect_obj_name));
+
+        // switch actor
+        if (cur_target->is_finished) {
+            switch_viewer_mode(FINISHED);
+        } else {
+            switch_viewer_mode(CONFIG);
+        }
+    }
+
+    auto calc_closest_point(const Eigen::Vector3d &point) {
+        Eigen::Vector3d close_point;
+        vtkIdType cell_id;
+        int sub_id;
+        double dis2;
+        cur_target->model_locator->FindClosestPoint(
+                point.data(), close_point.data(), cell_id, sub_id, dis2);
+        return std::make_tuple(close_point, std::sqrt(dis2));
+    }
+
+    // return if it needs to continue
+    bool calc_result() {
+        // prepare landmark
+        auto num_points = collected_points.size();
+        if (num_points < MIN_REG_POINTS) return true;
+        auto source_points = Eigen::Matrix3Xd{3, num_points};
+        auto target_points = Eigen::Matrix3Xd{3, num_points};
+        for (auto k = 0; k < num_points; ++k) {
+            source_points.col(k) = (*cur_target->all_points)[k];
+        }
+        for (auto k = 0; k < num_points; ++k) {
+            target_points.col(k) = collected_points[k];
+        }
+
+        // calculate landmark
+        auto result = (Eigen::Isometry3d) Eigen::umeyama(source_points, target_points, false);
+
+        // prepare icp
+        vtkNew<vtkPoints> icp_points;
+        auto landmark_inv = result.inverse();
+        for (auto k = 0; k < num_points; ++k) {
+            Eigen::Vector3d point = landmark_inv * collected_points[k];
+            icp_points->InsertNextPoint(point.data());
+        }
+
+        // calculate icp
+        vtkNew<vtkIterativeClosestPointTransform> icp;
+        vtkNew<vtkPolyData> tmp_poly;
+        tmp_poly->SetPoints(icp_points);
+        icp->GetLandmarkTransform()->SetModeToRigidBody();
+        icp->SetSource(tmp_poly);
+        icp->SetTarget(cur_target->model);
+        icp->Modified();
+        icp->Update();
+
+        // refine result
+        Eigen::Isometry3d trans_delta;
+        for (auto i = 0; i < 4; ++i)
+            for (auto j = 0; j < 4; ++j) {
+                trans_delta(i, j) = icp->GetMatrix()->GetElement(i, j);
+            }
+        result = result * trans_delta.inverse();
+
+        // commit result
+        conn->update_transform_variable(cur_target->target_var_name, result);
+
+        // calculate error only when all points are collected
+        if (num_points != cur_target->all_points->size()) return true;
+
+        // calculate error
+        cur_target->max_error = 0;
+        for (auto &points: cur_target->level_points) {
+            points->clear();
+        }
+        auto result_inv = result.inverse();
+        for (auto k = 0; k < num_points; ++k) {
+            Eigen::Vector3d point = result_inv * collected_points[k];
+
+            // find the closest point
+            double dis;
+            std::tie(std::ignore, dis) = calc_closest_point(point);
+
+            // update results
+            if (dis > cur_target->max_error) {
+                cur_target->max_error = dis;
+            }
+            if (dis < DIS_LIMIT_1) {
+                cur_target->level_points[0]->add_point(point);
+            } else if (dis < DIS_LIMIT_2) {
+                cur_target->level_points[1]->add_point(point);
+            } else {
+                cur_target->level_points[2]->add_point(point);
+            }
+        }
+
+        cur_target->is_finished = true;
+        return false;
+    }
+
+    void try_picking() {
+        assert(is_picking);
+        auto val = viewer->get_picked_point();
+        if (!val.has_value()) return;
+        cur_target->all_points->add_point(val.value());
+    }
+
+    void try_collect() {
+        assert(is_collecting);
+        auto val = conn->query_scalarxyz_variable(cur_target->point_var_name);
+        if (!val.has_value()) return;
+        collected_points.emplace_back(val.value());
+        SPDLOG_INFO("Collected point ({}, {}, {}).", val->x(), val->y(), val->z());
+        calc_result();
+        if (!progress_reg_point()) {
+            stop();
+        }
+    }
+
+    void show() {
+        if (ImGui::Begin("Registration Control")) {
+            if (is_collecting || is_picking) {
+                ImGui::BeginDisabled();
+            }
+            const char *target_name = (cur_target == nullptr) ? nullptr : cur_target->name.c_str();
+            if (ImGui::BeginCombo("Target", target_name)) {
+                for (auto k = 0; k < targets.size(); ++k) {
+                    bool is_selected = (k == cur_target_id);
+                    if (ImGui::Selectable(targets[k].name.c_str(), is_selected)) {
+                        cur_target_id = k;
+                        eq.emplace([this] { change_target(); });
+                    }
+                    if (is_selected) {
+                        ImGui::SetItemDefaultFocus();
+                    }
+                }
+                ImGui::EndCombo();
+            }
+            if (is_collecting || is_picking) {
+                ImGui::EndDisabled();
+            }
+            if (cur_target != nullptr) {
+                auto point_set = cur_target->all_points.get();
+                if (ImGui::CollapsingHeader("Actions")) {
+                    if (is_collecting) {
+                        ImGui::BeginDisabled();
+                    }
+                    if (ImGui::Checkbox("Config Points", &is_picking)) {
+                        if (is_picking) {
+                            eq.emplace([this] {
+                                switch_viewer_mode(CONFIG);
+                                viewer->start_picking();
+                            });
+                        } else {
+                            eq.emplace([this] { viewer->stop_picking(); });
+                        }
+                    }
+                    if (is_collecting) {
+                        ImGui::EndDisabled();
+                    }
+                    if (!is_picking && point_set->size() >= MIN_REG_POINTS) {
+                        ImGui::SameLine();
+                        if (!is_collecting) {
+                            if (ImGui::Button("Start")) {
+                                eq.emplace([this] { start(); });
+                            }
+                        } else {
+                            if (ImGui::Button("Stop")) {
+                                eq.emplace([this] { stop(); });
+                            }
+                        }
+                    }
+
+                }
+                if (ImGui::CollapsingHeader("Infos")) {
+                    void *token_delete = nullptr;
+                    if (ImGui::TreeNode("Fiducial Points")) {
+                        ImGui::PushItemWidth(200);
+                        point_set->for_each([&](void *token, const Eigen::Vector3d &point) {
+                            Eigen::Vector3f point_f = point.cast<float>();
+                            ImGui::PushID(token);
+                            ImGui::Bullet();
+                            ImGui::BeginDisabled();
+                            ImGui::InputFloat3("", point_f.data());
+                            ImGui::EndDisabled();
+                            if (is_picking) {
+                                ImGui::SameLine();
+                                if (ImGui::SmallButton("Delete")) {
+                                    token_delete = token;
+                                }
+                            }
+                            ImGui::PopID();
+                        });
+                        ImGui::PopItemWidth();
+                        ImGui::TreePop();
+                    }
+                    if (token_delete != nullptr) {
+                        eq.emplace([=] { point_set->remove_point(token_delete); });
+                    }
+                    ImGui::PushItemWidth(100);
+                    ImGui::Bullet();
+                    ImGui::InputFloat("Probe Test Error (mm)", &probe_test_error);
+                    if (cur_target->is_finished) {
+                        ImGui::Bullet();
+                        ImGui::InputFloat("Max Fiducial Error (mm)", &cur_target->max_error);
+                    }
+                    ImGui::PopItemWidth();
+                }
+            }
+        }
+        ImGui::End();
+        if (cur_target == nullptr) return;
+
+        if (ImGui::Begin("Registration View", nullptr, vtk_viewer::no_scroll_flag)) {
+            viewer->show();
+        }
+        ImGui::End();
+    }
+
+    void process() {
+        while (!eq.empty()) {
+            eq.front()();
+            eq.pop();
+        }
+        if (is_picking) {
+            try_picking();
+        } else if (is_collecting) {
+            try_collect();
+        }
+
+        // update probe transform
+        if (cur_target != nullptr) {
+            auto trans = conn->query_transform_variable(cur_target->probe_var_name);
+            update_actor_pose(probe_actor, trans);
+            if (trans.has_value()) {
+                Eigen::Vector3d point = trans.value().translation();
+                std::tie(std::ignore, probe_test_error) = calc_closest_point(point);
+            } else {
+                probe_test_error = NAN_VALUE;
+            }
+        }
+    }
+
+};
+
+registration::~registration() = default;
+
+registration *registration::create(const registration_config &conf) {
+    auto pimpl = new impl{};
+    pimpl->conn = conf.conn;
+    pimpl->probe_actor = create_actor(conf.probe_model_path);
+    auto ret = new registration{};
+    ret->pimpl.reset(pimpl);
+    return ret;
+}
+
+void registration::add_target(const registration_target &item) {
+    pimpl->add_target(item);
+}
+
+void registration::show() {
+    pimpl->show();
+}
+
+void registration::process() {
+    pimpl->process();
+}

src/components/registration.h

@@ -0,0 +1,48 @@
+#ifndef REMOTEAR3_REGISTRATION_H
+#define REMOTEAR3_REGISTRATION_H
+
+#include "core/local_connection.h"
+
+#include <vtkPolyData.h>
+
+#include <Eigen/Dense>
+
+#include <memory>
+#include <string_view>
+#include <vector>
+
+struct registration_config {
+    sophiar::local_connection *conn;
+    std::string probe_model_path;
+};
+
+struct registration_target {
+    std::string name;
+    std::string model_path;
+    std::vector<Eigen::Vector3d> pre_points;
+    std::string target_var_name;
+    std::string collect_var_name;
+    std::string collect_obj_name;
+    std::string probe_var_name;
+};
+
+class registration {
+public:
+
+    ~registration();
+
+    static registration *create(const registration_config &conf);
+
+    void add_target(const registration_target &item);
+
+    void show(); // show ImGui windows
+
+    void process();
+
+private:
+    struct impl;
+    std::unique_ptr<impl> pimpl;
+};
+
+
+#endif //REMOTEAR3_REGISTRATION_H

src/image_process/process_kernels.cu

@@ -238,7 +238,7 @@ struct enhance_image_func {
         using ImgElemT = decltype(in.x);
         static_assert(std::is_integral_v<ImgElemT>,
                       "Type of image element must be integer.");
-        ImgElemT c_maxgit  = max(max(in.x, in.y), in.z);
+        ImgElemT c_max  = max(max(in.x, in.y), in.z);
         ImgElemT c_min = min(min(in.x, in.y), in.z);
         ImgElemT delta = c_max - c_min;
 

src/main_ext.cpp

@@ -1,3 +1,4 @@
+#include "components/registration.h"
 #include "core/local_connection.h"
 #include "core/timestamp_helper.hpp"
 #include "cuda_helper.hpp"
@@ -30,7 +31,10 @@
 
 #include <boost/iostreams/device/mapped_file.hpp>
 
+#include <nlohmann/json.hpp>
+
 #include <cassert>
+#include <functional>
 #include <thread>
 #include <queue>
 #include <vector>
@@ -75,7 +79,12 @@ std::unique_ptr<frame_sender> sender;
 std::string sophiar_config_path;
 std::unique_ptr<std::thread> sophiar_thread;
 local_connection sophiar_conn;
-bool last_command_success = true;
+bool is_tracking = false;
+
+bool enable_reg = false;
+std::string probe_model_path;
+std::vector<registration_target> reg_targets;
+std::unique_ptr<registration> reg;
 
 bool show_vtk_debug = false;
 bool show_imgui_demo = false;
@@ -89,7 +98,7 @@ struct augment_store_type {
 
 std::vector<augment_store_type> augment_items;
 
-std::queue<void (*)()> simple_eq;
+std::queue<std::function<void()>> simple_eq;
 
 struct camera_related {
     uint64_t raw_cnt = 0;
@@ -132,6 +141,8 @@ struct camera_related {
                 glCopyImageSubData(augment_viewer->get_tex(), GL_TEXTURE_2D, 0, 0, 0, 0,
                                    augment_tex->id, GL_TEXTURE_2D, 0, 0, 0, 0, img_dev->cols, img_dev->rows, 1);
             }
+        } else {
+            augment_available = false;
         }
     }
 
@@ -235,13 +246,24 @@ void load_config() {
     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"]) {
         augment_items.emplace_back(
                 item["name"].as<std::string>(),
                 item["trans_var"].as<std::string>(),
-                load_stl(item["stl_file"].as<std::string>()));
+                create_actor(item["stl_file"].as<std::string>()));
+        if (auto reg_conf = item["registration"]; !reg_conf.IsNull()) {
+            reg_targets.emplace_back(
+                    item["name"].as<std::string>(),
+                    item["stl_file"].as<std::string>(),
+                    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
@@ -325,12 +347,20 @@ void initialize_main_window() {
     for (auto &item: augment_items) {
         augment_viewer->add_actor(item.actor);
     }
+    reg.reset(registration::create({&sophiar_conn,
+                                    probe_model_path}));
+    for (auto &item: reg_targets) {
+        reg->add_target(item);
+    }
 
     // initialize vtk test viewer
     vtk_test1 = std::make_unique<vtk_viewer>();
     vtk_test2 = std::make_unique<vtk_viewer>();
-    vtk_test1->add_actor(load_stl("/home/tpx/data/stls/GlassProbe_4Ball_3.STL"));
-    vtk_test2->add_actor(load_stl("/home/tpx/data/stls/HeadModelDrill.stl"));
+    vtk_test1->clear_actor();
+    vtk_test1->add_actor(create_actor("/home/tpx/data/stls/GlassProbe_4Ball_3.STL"));
+    vtk_test1->reset_camera();
+    vtk_test2->add_actor(create_actor("/home/tpx/data/stls/HeadModelDrill.stl"));
+    vtk_test1->start_picking();
 }
 
 bool is_camera_opened() {
@@ -506,7 +536,8 @@ void stop_encoder() {
 }
 
 void start_tracking() {
-    last_command_success = sophiar_conn.start_object("tracker_all");
+    CALL_CHECK(sophiar_conn.start_object("tracker_all"));
+    is_tracking = true;
 }
 
 void cleanup() {
@@ -549,20 +580,20 @@ void prepare_imgui_frame() {
             ImGui::SeparatorText("Actions");
             if (!is_camera_opened()) {
                 if (ImGui::Button("Open")) {
-                    simple_eq.push(open_cameras);
+                    simple_eq.emplace(open_cameras);
                 }
             } else { // cameras have been opened
                 if (ImGui::Button("Close")) {
-                    simple_eq.push(close_cameras);
+                    simple_eq.emplace(close_cameras);
                 }
                 ImGui::SameLine();
                 if (!is_capturing()) {
                     if (ImGui::Button("Start")) {
-                        simple_eq.push(start_capture);
+                        simple_eq.emplace(start_capture);
                     }
                 } else {
                     if (ImGui::Button("Stop")) {
-                        simple_eq.push(stop_capture);
+                        simple_eq.emplace(stop_capture);
                     }
 //                    if (!auto_save_raw) {
 //                        ImGui::SameLine();
@@ -576,17 +607,17 @@ void prepare_imgui_frame() {
             // camera configs
             ImGui::SeparatorText("Configs");
             if (ImGui::DragInt("Frame Rate (fps)", &capture_conf.frame_rate, 1, 1, 60)) {
-                simple_eq.push(upload_capture_config);
+                simple_eq.emplace(upload_capture_config);
                 main_encoder_conf.frame_rate = capture_conf.frame_rate;
-                simple_eq.push(upload_encoder_config);
+                simple_eq.emplace(upload_encoder_config);
             }
             if (ImGui::DragFloat("Exposure Time (ms)", &capture_conf.expo_time_ms,
                                  0.1, 0.1, 1e3f / capture_conf.frame_rate, "%.01f")) {
-                simple_eq.push(upload_capture_config);
+                simple_eq.emplace(upload_capture_config);
             }
             if (ImGui::DragFloat("Analog Gain (dB)", &capture_conf.gain_db,
                                  0.1, 0, 23.4, "%.01f")) {
-                simple_eq.push(upload_capture_config);
+                simple_eq.emplace(upload_capture_config);
             }
             ImGui::Checkbox("Enhance", &enhance_image);
 
@@ -643,18 +674,13 @@ void prepare_imgui_frame() {
 
             ImGui::SeparatorText("Actions");
             if (ImGui::Button("Start Tracking")) {
-                simple_eq.push(start_tracking);
+                simple_eq.emplace(start_tracking);
             }
-
-            ImGui::SeparatorText("Infos");
-            if (last_command_success) {
-                ImGui::PushStyleColor(ImGuiCol_Text, (ImVec4) ImColor(0, 255, 0));
-            } else {
-                ImGui::PushStyleColor(ImGuiCol_Text, (ImVec4) ImColor(255, 0, 0));
+            if (is_tracking) {
+                ImGui::Checkbox("Registration Panel", &enable_reg);
             }
-            ImGui::LabelText("Last Command", last_command_success ? "Success" : "Failed");
-            ImGui::PopStyleColor();
 
+            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();
@@ -670,10 +696,10 @@ void prepare_imgui_frame() {
                 }
             };
 //            ImGui::BeginDisabled();
-            helper_func("camera_ref_in_tracker", "Camera");
+            helper_func("camera_ref_in_tracker", "Camera"); // TODO set in config
             helper_func("probe_in_tracker", "Probe");
-            helper_func("femur_in_tracker", "Femur");
-            helper_func("tibia_in_tracker", "Tibia", true);
+            helper_func("femur_ref_in_tracker", "Femur");
+            helper_func("tibia_ref_in_tracker", "Tibia", true);
 //            ImGui::EndDisabled();
 
             ImGui::PopID();
@@ -704,18 +730,18 @@ void prepare_imgui_frame() {
             ImGui::SeparatorText("Actions");
             if (!is_encoding()) {
                 if (ImGui::Button("Start")) {
-                    simple_eq.push(start_encoder);
+                    simple_eq.emplace(start_encoder);
                 }
             } else {
                 if (ImGui::Button("Close")) {
-                    simple_eq.push(stop_encoder);
+                    simple_eq.emplace(stop_encoder);
                 }
             }
 
             ImGui::SeparatorText("Configs");
             if (ImGui::DragFloat("Bitrate (Mbps)", &main_encoder_conf.bitrate_mbps,
                                  0.1, 1, 20, "%.01f")) {
-                simple_eq.push(upload_encoder_config);
+                simple_eq.emplace(upload_encoder_config);
             }
 
             if (is_encoding()) {
@@ -742,11 +768,11 @@ void prepare_imgui_frame() {
             ImGui::SeparatorText("Actions");
             if (!is_sending()) {
                 if (ImGui::Button("Start")) {
-                    simple_eq.push(start_sender);
+                    simple_eq.emplace(start_sender);
                 }
             } else {
                 if (ImGui::Button("Stop")) {
-                    simple_eq.push(stop_sender);
+                    simple_eq.emplace(stop_sender);
                 }
             }
 
@@ -785,6 +811,10 @@ void prepare_imgui_frame() {
         ImGui::End();
     }
 
+    if (enable_reg) {
+        reg->show();
+    }
+
     ImGui::Render();
 }
 
@@ -793,6 +823,11 @@ void handle_imgui_events() {
         simple_eq.front()();
         simple_eq.pop();
     }
+
+    // registration related
+    if (enable_reg) {
+        reg->process();
+    }
 }
 
 void wait_camera_frames() {
@@ -805,9 +840,7 @@ void process_camera_frames() {
     // update augment items' position
     for (auto &item: augment_items) {
         auto trans = sophiar_conn.query_transform_variable(item.trans_var);
-        if (trans.has_value()) {
-            update_actor_pose(item.actor, trans.value());
-        }
+        update_actor_pose(item.actor, trans);
     }
 
     left.process_frame(IMG_RAW_HOST_LEFT);

src/vtk_viewer.cpp

@@ -1,30 +1,88 @@
 #include "vtk_viewer.h"
 
 #include <vtkCamera.h>
+#include <vtkCellArray.h>
 #include <vtkGenericOpenGLRenderWindow.h>
 #include <vtkGenericRenderWindowInteractor.h>
 #include <vtkInteractorStyleTrackballCamera.h>
 #include <vtkMatrix4x4.h>
 #include <vtkNew.h>
 #include <vtkOpenGLFramebufferObject.h>
+#include <vtkPointPicker.h>
+#include <vtkPoints.h>
 #include <vtkPolyDataMapper.h>
+#include <vtkProperty.h>
 #include <vtkRenderer.h>
 #include <vtkSTLReader.h>
 #include <vtkTextureObject.h>
 
 #include <spdlog/spdlog.h>
 
+#include <vector>
+
 struct vtk_viewer::impl {
 
+    struct versatile_interaction_style;
+
     static constexpr auto default_focal_length = 8; // 8mm
 
     cv::Size last_size;
     vtkSmartPointer<vtkGenericOpenGLRenderWindow> window;
     vtkSmartPointer<vtkCamera> camera;
     vtkSmartPointer<vtkRenderWindowInteractor> controller;
-    vtkSmartPointer<vtkInteractorStyle> style;
+    vtkSmartPointer<versatile_interaction_style> style;
     vtkSmartPointer<vtkRenderer> renderer;
 
+    bool is_picking = false;
+    std::optional<Eigen::Vector3d> picked_point;
+
+    struct versatile_interaction_style
+            : public vtkInteractorStyleTrackballCamera {
+    vtkTypeMacro(versatile_interaction_style, vtkInteractorStyleTrackballCamera);
+
+        static versatile_interaction_style *New() {
+            auto ret = new versatile_interaction_style{};
+            ret->InitializeObjectBase();
+            return ret;
+        }
+
+        impl *pimpl = nullptr;
+        Eigen::Vector2i mouse_down_pose;
+
+        Eigen::Vector2i get_mouse_pos() {
+            return {Interactor->GetEventPosition()[0],
+                    Interactor->GetEventPosition()[1],
+            };
+        }
+
+        void OnLeftButtonDown() override {
+            mouse_down_pose = get_mouse_pos();
+            vtkInteractorStyleTrackballCamera::OnLeftButtonDown();
+        }
+
+        bool try_picking() {
+            assert(pimpl != nullptr);
+            if (!pimpl->is_picking) return false;
+            auto ret = Interactor->GetPicker()->Pick(mouse_down_pose.x(),
+                                                     mouse_down_pose.y(),
+                                                     0,
+                                                     pimpl->renderer);
+            if (ret == 0) return true; // does not pick anything
+            Eigen::Vector3d point;
+            Interactor->GetPicker()->GetPickPosition(point.data());
+            pimpl->picked_point = point;
+            SPDLOG_INFO("Picked point ({}, {}, {}).", point.x(), point.y(), point.z());
+            return true;
+        }
+
+        void OnLeftButtonUp() override {
+            if (get_mouse_pos() == mouse_down_pose) {
+                try_picking();
+            }
+            vtkInteractorStyleTrackballCamera::OnLeftButtonUp();
+        }
+    };
+
     static cv::Size to_cv_size(ImVec2 size) {
         return {(int) size.x, (int) size.y};
     }
@@ -40,14 +98,15 @@ struct vtk_viewer::impl {
 //        window->SetFrameBlitModeToNoBlit();
 
         camera = vtkSmartPointer<vtkCamera>::New();
-        camera->SetClippingRange(100, 2000); // 10cm to 2m
+        camera->SetClippingRange(default_focal_length, 2000); // 8mm to 2m
 
         renderer = vtkSmartPointer<vtkRenderer>::New();
         renderer->SetUseDepthPeeling(true);
         renderer->SetActiveCamera(camera);
         window->AddRenderer(renderer);
 
-        style = vtkSmartPointer<vtkInteractorStyleTrackballCamera>::New();
+        style = vtkSmartPointer<versatile_interaction_style>::New();
+        style->pimpl = this;
         style->SetDefaultRenderer(renderer);
         controller = vtkSmartPointer<vtkGenericRenderWindowInteractor>::New();
         controller->SetInteractorStyle(style);
@@ -71,6 +130,7 @@ struct vtk_viewer::impl {
         if (size != last_size) [[unlikely]] {
             controller->SetSize(size.width, size.height);
             window->SetSize(size.width, size.height);
+            last_size = size;
         }
         if (interactive) {
             process_event();
@@ -84,7 +144,7 @@ struct vtk_viewer::impl {
         // set event position
         auto &io = ImGui::GetIO();
         io.ConfigWindowsMoveFromTitleBarOnly = true;
-        auto view_pos = ImGui::GetCursorStartPos();
+        auto view_pos = ImGui::GetWindowPos();
         auto x_pos = io.MousePos.x - view_pos.x;
         auto y_pos = io.MousePos.y - view_pos.y;
         controller->SetEventInformationFlipY(x_pos, y_pos, io.KeyCtrl, io.KeyShift);
@@ -93,7 +153,14 @@ struct vtk_viewer::impl {
         if (ImGui::IsWindowHovered()) {
             if (io.MouseClicked[ImGuiMouseButton_Left]) {
                 controller->InvokeEvent(vtkCommand::LeftButtonPressEvent);
-            } else if (io.MouseWheel > 0) {
+            }
+            if (io.MouseClicked[ImGuiMouseButton_Middle]) {
+                controller->InvokeEvent(vtkCommand::MiddleButtonPressEvent);
+            }
+            if (io.MouseClicked[ImGuiMouseButton_Right]) {
+                controller->InvokeEvent(vtkCommand::RightButtonPressEvent);
+            }
+            if (io.MouseWheel > 0) {
                 controller->InvokeEvent(vtkCommand::MouseWheelForwardEvent);
             } else if (io.MouseWheel < 0) {
                 controller->InvokeEvent(vtkCommand::MouseWheelBackwardEvent);
@@ -102,6 +169,12 @@ struct vtk_viewer::impl {
         if (io.MouseReleased[ImGuiMouseButton_Left]) {
             controller->InvokeEvent(vtkCommand::LeftButtonReleaseEvent);
         }
+        if (io.MouseReleased[ImGuiMouseButton_Middle]) {
+            controller->InvokeEvent(vtkCommand::MiddleButtonReleaseEvent);
+        }
+        if (io.MouseReleased[ImGuiMouseButton_Right]) {
+            controller->InvokeEvent(vtkCommand::RightButtonReleaseEvent);
+        }
         controller->InvokeEvent(vtkCommand::MouseMoveEvent);
     }
 
@@ -137,13 +210,16 @@ GLuint vtk_viewer::get_tex() const {
 
 void vtk_viewer::add_actor(vtkActor *actor) {
     pimpl->renderer->AddActor(actor);
-    pimpl->renderer->ResetCamera();
 }
 
 void vtk_viewer::remove_actor(vtkActor *actor) {
     pimpl->renderer->RemoveActor(actor);
 }
 
+void vtk_viewer::clear_actor() {
+    pimpl->renderer->RemoveAllViewProps();
+}
+
 void vtk_viewer::show(const std::string &name) {
     pimpl->show_imgui_window(name.c_str(), ImGui::GetContentRegionAvail());
 }
@@ -157,12 +233,38 @@ void vtk_viewer::set_camera_view_angle(double angle) {
     pimpl->camera->Modified();
 }
 
-vtkSmartPointer<vtkActor> load_stl(const std::string &path) {
+void vtk_viewer::start_picking() {
+    pimpl->is_picking = true;
+}
+
+void vtk_viewer::stop_picking() {
+    pimpl->is_picking = false;
+}
+
+bool vtk_viewer::is_picking() {
+    return pimpl->is_picking;
+}
+
+std::optional<Eigen::Vector3d> vtk_viewer::get_picked_point() {
+    auto ret = pimpl->picked_point;
+    pimpl->picked_point = {};
+    return ret;
+}
+
+void vtk_viewer::reset_camera() {
+    pimpl->renderer->ResetCamera();
+}
+
+vtkSmartPointer<vtkPolyData> load_stl(const std::string &path) {
     vtkNew<vtkSTLReader> reader;
     reader->SetFileName(path.c_str());
     reader->Update();
+    return reader->GetOutput();
+}
+
+vtkSmartPointer<vtkActor> create_actor(vtkPolyData *data) {
     vtkNew<vtkPolyDataMapper> mapper;
-    mapper->SetInputData(reader->GetOutput());
+    mapper->SetInputData(data);
     vtkNew<vtkActor> actor;
     actor->SetMapper(mapper);
     vtkNew<vtkMatrix4x4> pose;
@@ -170,14 +272,139 @@ vtkSmartPointer<vtkActor> load_stl(const std::string &path) {
     return actor;
 }
 
-void update_actor_pose(vtkActor *actor, const Eigen::Isometry3d &trans) {
+vtkSmartPointer<vtkActor> create_actor(const std::string &path) {
+    return create_actor(load_stl(path));
+}
+
+void update_actor_pose(vtkActor *actor, const std::optional<Eigen::Isometry3d> &trans) {
+    if (!trans.has_value()) {
+        actor->VisibilityOff();
+        return;
+    }
+    actor->VisibilityOn();
+    auto &real_trans = trans.value();
     auto matrix = actor->GetUserMatrix();
     assert(matrix != nullptr);
     for (int i = 0; i < 4; ++i) {
         for (int j = 0; j < 4; ++j) {
-            matrix->SetElement(i, j, trans(i, j));
+            matrix->SetElement(i, j, real_trans(i, j));
         }
     }
     matrix->Modified();
     actor->Modified();
+}
+
+struct smart_point_sets::impl {
+
+    vtkSmartPointer<vtkPoints> points;
+    vtkSmartPointer<vtkCellArray> vertices;
+    vtkSmartPointer<vtkPolyData> poly;
+    vtkSmartPointer<vtkActor> actor;
+
+    std::vector<Eigen::Vector3d> points_store;
+
+    impl() {
+        actor = vtkSmartPointer<vtkActor>::New();
+        actor->GetProperty()->SetRenderPointsAsSpheres(true);
+        actor->GetProperty()->SetPointSize(10);
+        reconstruct();
+    }
+
+    void add_point_to_poly(const Eigen::Vector3d &point) {
+        vtkIdType pid[1];
+        pid[0] = points->InsertNextPoint(point.data());
+        points->Modified();
+        vertices->InsertNextCell(1, pid);
+        vertices->Modified();
+        poly->Modified();
+    }
+
+    auto add_point(const Eigen::Vector3d &point) {
+        points_store.emplace_back(point);
+        add_point_to_poly(point);
+    }
+
+    void reconstruct() {
+        points = vtkSmartPointer<vtkPoints>::New();
+        vertices = vtkSmartPointer<vtkCellArray>::New();
+        poly = vtkSmartPointer<vtkPolyData>::New();
+        poly->SetPoints(points);
+        poly->SetVerts(vertices);
+        vtkNew<vtkPolyDataMapper> mapper;
+        mapper->SetInputData(poly);
+        actor->SetMapper(mapper);
+        actor->Modified();
+
+        for (auto &p: points_store) {
+            add_point_to_poly(p);
+        }
+    }
+
+    void for_each(const for_each_func_type &func) {
+        for (auto iter = points_store.begin();
+             iter != points_store.end();
+             ++iter) {
+            func(*(void **) &iter, *iter);
+        }
+    }
+
+    void remove_point(void *token) {
+        using iter_type = decltype(points_store)::iterator;
+        points_store.erase(*(iter_type *) (&token));
+        reconstruct();
+    }
+
+    Eigen::Vector3d pop_front() {
+        assert(!points_store.empty());
+        auto ret = points_store.front();
+        points_store.erase(points_store.begin());
+        reconstruct();
+        return ret;
+    }
+};
+
+smart_point_sets::smart_point_sets()
+        : pimpl(std::make_unique<impl>()) {}
+
+smart_point_sets::~smart_point_sets() = default;
+
+void smart_point_sets::add_point(const Eigen::Vector3d &point) {
+    pimpl->add_point(point);
+}
+
+vtkActor *smart_point_sets::get_actor() {
+    return pimpl->actor;
+}
+
+void smart_point_sets::for_each(const for_each_func_type &func) {
+    pimpl->for_each(func);
+}
+
+void smart_point_sets::remove_point(void *token) {
+    pimpl->remove_point(token);
+}
+
+Eigen::Vector3d smart_point_sets::pop_front() {
+    return pimpl->pop_front();
+}
+
+bool smart_point_sets::empty() const {
+    return pimpl->points_store.empty();
+}
+
+size_t smart_point_sets::size() const {
+    return pimpl->points_store.size();
+}
+
+std::vector<Eigen::Vector3d> smart_point_sets::to_vector() const {
+    return pimpl->points_store;
+}
+
+Eigen::Vector3d smart_point_sets::operator[](size_t index) const {
+    return pimpl->points_store[index];
+}
+
+void smart_point_sets::clear() const {
+    pimpl->points_store.clear();
+    pimpl->reconstruct();
 }

src/vtk_viewer.h

@@ -10,13 +10,20 @@
 #include <Eigen/Geometry>
 
 #include <vtkActor.h>
+#include <vtkPolyData.h>
 #include <vtkSmartPointer.h>
 
+#include <functional>
 #include <memory>
+#include <optional>
 
-vtkSmartPointer<vtkActor> load_stl(const std::string &path);
+vtkSmartPointer<vtkPolyData> load_stl(const std::string &path);
 
-void update_actor_pose(vtkActor *actor, const Eigen::Isometry3d &trans);
+vtkSmartPointer<vtkActor> create_actor(vtkPolyData *data);
+
+vtkSmartPointer<vtkActor> create_actor(const std::string &path);
+
+void update_actor_pose(vtkActor *actor, const std::optional<Eigen::Isometry3d> &trans);
 
 class vtk_viewer {
 public:
@@ -32,20 +39,66 @@ public:
 
     void remove_actor(vtkActor *actor);
 
+    void clear_actor();
+
     void set_camera_view_angle(double angle);
 
     void set_camera_pose(const Eigen::Isometry3d &trans);
 
+    void reset_camera();
+
     void render(cv::Size size);
 
     GLuint get_tex() const;
 
     void show(const std::string &name = "VtkViewer"); // Show in ImGui as child
 
+    // picking related
+    void start_picking();
+
+    void stop_picking();
+
+    bool is_picking();
+
+    std::optional<Eigen::Vector3d> get_picked_point();
+
 private:
     struct impl;
     std::unique_ptr<impl> pimpl;
 };
 
+class smart_point_sets {
+public:
+
+    smart_point_sets();
+
+    ~smart_point_sets();
+
+    vtkActor *get_actor();
+
+    void add_point(const Eigen::Vector3d &point);
+
+    void remove_point(void *token);
+
+    size_t size() const;
+
+    bool empty() const;
+
+    void clear() const;
+
+    Eigen::Vector3d pop_front();
+
+    using for_each_func_type = std::function<void(void *, Eigen::Vector3d)>; // void * means a token
+
+    void for_each(const for_each_func_type &func);
+
+    std::vector<Eigen::Vector3d> to_vector() const;
+
+    Eigen::Vector3d operator[](size_t index) const;
+
+private:
+    struct impl;
+    std::unique_ptr<impl> pimpl;
+};
 
 #endif //REMOTEAR3_VTK_VIEWER_H