Restructuring some code.

extras/ros/README.md

@@ -0,0 +1,20 @@
+# ROS-Wrapper
+
+Run pose estimator with ros topics as inputs and publish detected poses.
+
+<br>
+
+- Build container:
+  ```bash
+  docker build --progress=plain -t rapidposetriangulation_ros -f extras/ros/dockerfile .
+  ```
+
+- Update or remove the `ROS_DOMAIN_ID` in `docker-compose.yml` to fit your environment
+
+- Run and test:
+  ```bash
+  xhost +; docker compose -f extras/ros/docker-compose.yml up
+
+  docker exec -it ros-test_node-1 bash
+  export ROS_DOMAIN_ID=18
+  ```

extras/ros/docker-compose.yml

@@ -10,8 +10,8 @@ services:
         runtime: nvidia
         privileged: true
         volumes:
-            - ../:/RapidPoseTriangulation/
+            - ../../:/RapidPoseTriangulation/
-            - ../skelda/:/skelda/
+            - ../../skelda/:/skelda/
             - /tmp/.X11-unix:/tmp/.X11-unix
             - /dev/shm:/dev/shm
         environment:
@@ -27,8 +27,8 @@ services:
         runtime: nvidia
         privileged: true
         volumes:
-            - ../:/RapidPoseTriangulation/
+            - ../../:/RapidPoseTriangulation/
-            - ../skelda/:/skelda/
+            - ../../skelda/:/skelda/
             - /tmp/.X11-unix:/tmp/.X11-unix
             - /dev/shm:/dev/shm
         environment:
@@ -45,8 +45,8 @@ services:
         runtime: nvidia
         privileged: true
         volumes:
-            - ../:/RapidPoseTriangulation/
+            - ../../:/RapidPoseTriangulation/
-            - ../skelda/:/skelda/
+            - ../../skelda/:/skelda/
             - /tmp/.X11-unix:/tmp/.X11-unix
             - /dev/shm:/dev/shm
         environment:

extras/ros/dockerfile

@@ -60,14 +60,16 @@ RUN echo "source /project/dev_ws/install/setup.bash" >> ~/.bashrc
 RUN pip3 install --no-cache-dir "setuptools<=58.2.0"
 
 # Copy modules
-COPY ./ros/pose2D_visualizer /RapidPoseTriangulation/ros/pose2D_visualizer/
+COPY ./extras/include/ /RapidPoseTriangulation/extras/include/
-COPY ./ros/rpt2D_wrapper_py /RapidPoseTriangulation/ros/rpt2D_wrapper_py/
+COPY ./scripts/ /RapidPoseTriangulation/scripts/
-COPY ./ros/rpt2D_wrapper_cpp /RapidPoseTriangulation/ros/rpt2D_wrapper_cpp/
+COPY ./extras/ros/pose2D_visualizer/ /RapidPoseTriangulation/extras/ros/pose2D_visualizer/
+COPY ./extras/ros/rpt2D_wrapper_py/ /RapidPoseTriangulation/extras/ros/rpt2D_wrapper_py/
+COPY ./extras/ros/rpt2D_wrapper_cpp/ /RapidPoseTriangulation/extras/ros/rpt2D_wrapper_cpp/
 
 # Link and build as ros package
-RUN ln -s  /RapidPoseTriangulation/ros/pose2D_visualizer/ /project/dev_ws/src/pose2D_visualizer
+RUN ln -s  /RapidPoseTriangulation/extras/ros/pose2D_visualizer/ /project/dev_ws/src/pose2D_visualizer
-RUN ln -s  /RapidPoseTriangulation/ros/rpt2D_wrapper_py/ /project/dev_ws/src/rpt2D_wrapper_py
+RUN ln -s  /RapidPoseTriangulation/extras/ros/rpt2D_wrapper_py/ /project/dev_ws/src/rpt2D_wrapper_py
-RUN ln -s  /RapidPoseTriangulation/ros/rpt2D_wrapper_cpp/ /project/dev_ws/src/rpt2D_wrapper_cpp
+RUN ln -s  /RapidPoseTriangulation/extras/ros/rpt2D_wrapper_cpp/ /project/dev_ws/src/rpt2D_wrapper_cpp
 RUN /bin/bash -i -c 'cd /project/dev_ws/; colcon build --symlink-install --cmake-args -DCMAKE_BUILD_TYPE=Release'
 
 # Update ros packages -> autocompletion and check

extras/ros/pose2D_visualizer/pose2D_visualizer/pose2D_visualizer.py

@@ -14,8 +14,8 @@ from sensor_msgs.msg import Image
 from std_msgs.msg import String
 
 filepath = os.path.dirname(os.path.realpath(__file__)) + "/"
-sys.path.append(filepath + "../../../scripts/")
+sys.path.append(filepath + "../../../../scripts/")
-import test_triangulate
+import utils_pipeline
 
 from skelda import utils_view
 
@@ -42,7 +42,7 @@ def callback_images(image_data):
     # Convert into cv images from image string
     if image_data.encoding == "bayer_rggb8":
         bayer_image = bridge.imgmsg_to_cv2(image_data, "bayer_rggb8")
-        color_image = test_triangulate.bayer2rgb(bayer_image)
+        color_image = utils_pipeline.bayer2rgb(bayer_image)
     elif image_data.encoding == "mono8":
         gray_image = bridge.imgmsg_to_cv2(image_data, "mono8")
         color_image = cv2.cvtColor(gray_image, cv2.COLOR_GRAY2RGB)

extras/ros/rpt2D_wrapper_cpp/CMakeLists.txt

@@ -23,11 +23,11 @@ find_package(sensor_msgs REQUIRED)
 find_package(cv_bridge REQUIRED)
 find_package(OpenCV REQUIRED)
 
-### 3) ONNX Runtime 
+### 3) ONNX Runtime
-include_directories(/onnxruntime/include
+include_directories(/onnxruntime/include/
-                    /onnxruntime/include/onnxruntime/core/session
+                    /onnxruntime/include/onnxruntime/core/session/
-                    /onnxruntime/include/onnxruntime/core/providers/tensorrt)
+                    /onnxruntime/include/onnxruntime/core/providers/tensorrt/)
-link_directories(/onnxruntime/build/Linux/Release)
+link_directories(/onnxruntime/build/Linux/Release/)
 
 add_executable(rpt2D_wrapper src/rpt2D_wrapper.cpp)
 ament_target_dependencies(rpt2D_wrapper rclcpp std_msgs sensor_msgs cv_bridge)

extras/ros/rpt2D_wrapper_cpp/src/rpt2D_wrapper.cpp

@@ -18,11 +18,11 @@
 #include <cv_bridge/cv_bridge.h>
 
 // JSON library
-#include "nlohmann/json.hpp"
+#include "/RapidPoseTriangulation/extras/include/nlohmann/json.hpp"
 using json = nlohmann::json;
 
-#include "test_triangulate.hpp"
+#include "/RapidPoseTriangulation/scripts/utils_pipeline.hpp"
-#include "utils_2d_pose.hpp"
+#include "/RapidPoseTriangulation/scripts/utils_2d_pose.hpp"
 
 // =================================================================================================
 
@@ -34,6 +34,12 @@ static const float min_bbox_score = 0.4;
 static const float min_bbox_area = 0.1 * 0.1;
 static const bool batch_poses = true;
 
+static const std::map<std::string, bool> whole_body = {
+    {"foots", true},
+    {"face", true},
+    {"hands", true},
+};
+
 // =================================================================================================
 // =================================================================================================
 
@@ -61,9 +67,9 @@ public:
         pose_pub_ = this->create_publisher<std_msgs::msg::String>(pose_out_topic, qos_profile);
 
         // Load model
-        bool whole_body = test_triangulate::use_whole_body();
+        bool use_wb = utils_pipeline::use_whole_body(whole_body);
         this->kpt_model = std::make_unique<utils_2d_pose::PosePredictor>(
-            whole_body, min_bbox_score, min_bbox_area, batch_poses);
+            use_wb, min_bbox_score, min_bbox_area, batch_poses);
 
         RCLCPP_INFO(this->get_logger(), "Finished initialization of pose estimator.");
 
@@ -86,6 +92,7 @@ private:
 
     // Pose model pointer
     std::unique_ptr<utils_2d_pose::PosePredictor> kpt_model;
+    const std::vector<std::string> joint_names_2d = utils_pipeline::get_joint_names(whole_body);
 
     // Threading
     std::thread model_thread;
@@ -138,7 +145,7 @@ void Rpt2DWrapperNode::callbackImages(const sensor_msgs::msg::Image::SharedPtr m
         {
             cv_bridge::CvImageConstPtr cv_ptr = cv_bridge::toCvShare(msg, "rgb8");
             cv::Mat color_image = cv_ptr->image;
-            bayer_image = test_triangulate::rgb2bayer(color_image);
+            bayer_image = utils_pipeline::rgb2bayer(color_image);
         }
         else
         {
@@ -185,14 +192,14 @@ void Rpt2DWrapperNode::callbackModel()
     }
 
     // Create image vector
-    cv::Mat rgb_image = test_triangulate::bayer2rgb(local_image);
+    cv::Mat rgb_image = utils_pipeline::bayer2rgb(local_image);
     std::vector<cv::Mat> images_2d;
     images_2d.push_back(rgb_image);
 
     // Predict 2D poses
     auto poses_2d_all = kpt_model->predict(images_2d);
-    auto poses_2d_upd = test_triangulate::update_keypoints(
+    auto poses_2d_upd = utils_pipeline::update_keypoints(
-        poses_2d_all, test_triangulate::joint_names_2d);
+        poses_2d_all, joint_names_2d, whole_body);
     auto &poses_2d = poses_2d_upd[0];
 
     // Drop persons with no joints
@@ -228,7 +235,7 @@ void Rpt2DWrapperNode::callbackModel()
 
     json poses_msg;
     poses_msg["bodies2D"] = valid_poses; // shape: persons x keypoints x 3
-    poses_msg["joints"] = test_triangulate::joint_names_2d;
+    poses_msg["joints"] = joint_names_2d;
     poses_msg["num_persons"] = valid_poses.size();
     poses_msg["timestamps"] = {
         {"image", local_timestamp},

ros/rpt2D_wrapper_cpp/src/test_triangulate.hpp

@@ -1,296 +0,0 @@
-#pragma once
-
-#include <string>
-#include <vector>
-
-#include <opencv2/opencv.hpp>
-
-// =================================================================================================
-
-namespace test_triangulate
-{
-    // If any part shall be disabled, also update the joint names list
-    static const std::map<std::string, bool> whole_body = {
-        {"foots", true},
-        {"face", true},
-        {"hands", true},
-    };
-
-    static const std::vector<std::string> joint_names_2d = {
-        // coco
-        "nose",
-        "eye_left",
-        "eye_right",
-        "ear_left",
-        "ear_right",
-        "shoulder_left",
-        "shoulder_right",
-        "elbow_left",
-        "elbow_right",
-        "wrist_left",
-        "wrist_right",
-        "hip_left",
-        "hip_right",
-        "knee_left",
-        "knee_right",
-        "ankle_left",
-        "ankle_right",
-        // foot
-        "foot_toe_big_left",
-        "foot_toe_small_left",
-        "foot_heel_left",
-        "foot_toe_big_right",
-        "foot_toe_small_right",
-        "foot_heel_right",
-        // face
-        "face_jaw_right_1",
-        "face_jaw_right_2",
-        "face_jaw_right_3",
-        "face_jaw_right_4",
-        "face_jaw_right_5",
-        "face_jaw_right_6",
-        "face_jaw_right_7",
-        "face_jaw_right_8",
-        "face_jaw_middle",
-        "face_jaw_left_1",
-        "face_jaw_left_2",
-        "face_jaw_left_3",
-        "face_jaw_left_4",
-        "face_jaw_left_5",
-        "face_jaw_left_6",
-        "face_jaw_left_7",
-        "face_jaw_left_8",
-        "face_eyebrow_right_1",
-        "face_eyebrow_right_2",
-        "face_eyebrow_right_3",
-        "face_eyebrow_right_4",
-        "face_eyebrow_right_5",
-        "face_eyebrow_left_1",
-        "face_eyebrow_left_2",
-        "face_eyebrow_left_3",
-        "face_eyebrow_left_4",
-        "face_eyebrow_left_5",
-        "face_nose_1",
-        "face_nose_2",
-        "face_nose_3",
-        "face_nose_4",
-        "face_nose_5",
-        "face_nose_6",
-        "face_nose_7",
-        "face_nose_8",
-        "face_nose_9",
-        "face_eye_right_1",
-        "face_eye_right_2",
-        "face_eye_right_3",
-        "face_eye_right_4",
-        "face_eye_right_5",
-        "face_eye_right_6",
-        "face_eye_left_1",
-        "face_eye_left_2",
-        "face_eye_left_3",
-        "face_eye_left_4",
-        "face_eye_left_5",
-        "face_eye_left_6",
-        "face_mouth_1",
-        "face_mouth_2",
-        "face_mouth_3",
-        "face_mouth_4",
-        "face_mouth_5",
-        "face_mouth_6",
-        "face_mouth_7",
-        "face_mouth_8",
-        "face_mouth_9",
-        "face_mouth_10",
-        "face_mouth_11",
-        "face_mouth_12",
-        "face_mouth_13",
-        "face_mouth_14",
-        "face_mouth_15",
-        "face_mouth_16",
-        "face_mouth_17",
-        "face_mouth_18",
-        "face_mouth_19",
-        "face_mouth_20",
-        // hand
-        "hand_wrist_left",
-        "hand_finger_thumb_left_1",
-        "hand_finger_thumb_left_2",
-        "hand_finger_thumb_left_3",
-        "hand_finger_thumb_left_4",
-        "hand_finger_index_left_1",
-        "hand_finger_index_left_2",
-        "hand_finger_index_left_3",
-        "hand_finger_index_left_4",
-        "hand_finger_middle_left_1",
-        "hand_finger_middle_left_2",
-        "hand_finger_middle_left_3",
-        "hand_finger_middle_left_4",
-        "hand_finger_ring_left_1",
-        "hand_finger_ring_left_2",
-        "hand_finger_ring_left_3",
-        "hand_finger_ring_left_4",
-        "hand_finger_pinky_left_1",
-        "hand_finger_pinky_left_2",
-        "hand_finger_pinky_left_3",
-        "hand_finger_pinky_left_4",
-        "hand_wrist_right",
-        "hand_finger_thumb_right_1",
-        "hand_finger_thumb_right_2",
-        "hand_finger_thumb_right_3",
-        "hand_finger_thumb_right_4",
-        "hand_finger_index_right_1",
-        "hand_finger_index_right_2",
-        "hand_finger_index_right_3",
-        "hand_finger_index_right_4",
-        "hand_finger_middle_right_1",
-        "hand_finger_middle_right_2",
-        "hand_finger_middle_right_3",
-        "hand_finger_middle_right_4",
-        "hand_finger_ring_right_1",
-        "hand_finger_ring_right_2",
-        "hand_finger_ring_right_3",
-        "hand_finger_ring_right_4",
-        "hand_finger_pinky_right_1",
-        "hand_finger_pinky_right_2",
-        "hand_finger_pinky_right_3",
-        "hand_finger_pinky_right_4",
-        // extras
-        "hip_middle",
-        "shoulder_middle",
-        "head",
-    };
-
-    // =============================================================================================
-    // =============================================================================================
-
-    [[maybe_unused]] static bool use_whole_body()
-    {
-        for (const auto &pair : whole_body)
-        {
-            if (pair.second)
-            {
-                return true;
-            }
-        }
-        return false;
-    }
-
-    // =============================================================================================
-
-    cv::Mat bayer2rgb(const cv::Mat &bayer)
-    {
-        cv::Mat rgb;
-        cv::cvtColor(bayer, rgb, cv::COLOR_BayerBG2RGB);
-        return rgb;
-    }
-
-    cv::Mat rgb2bayer(const cv::Mat &img)
-    {
-        CV_Assert(img.type() == CV_8UC3);
-        cv::Mat bayer(img.rows, img.cols, CV_8UC1);
-
-        for (int r = 0; r < img.rows; ++r)
-        {
-            const uchar *imgData = img.ptr<uchar>(r);
-            uchar *bayerRowPtr = bayer.ptr<uchar>(r);
-
-            for (int c = 0; c < img.cols; ++c)
-            {
-                int pixelIndex = 3 * c;
-
-                // Use faster bit operation instead of modulo+if
-                // Even row, even col => R = 0
-                // Even row, odd col => G = 1
-                // Odd row, even col => G = 1
-                // Odd row, odd col => B = 2
-                int row_mod = r & 1;
-                int col_mod = c & 1;
-                int component = row_mod + col_mod;
-
-                bayerRowPtr[c] = imgData[pixelIndex + component];
-            }
-        }
-
-        return bayer;
-    }
-
-    // =============================================================================================
-
-    inline int find_index(const std::vector<std::string> &vec, const std::string &key)
-    {
-        auto it = std::find(vec.begin(), vec.end(), key);
-        if (it == vec.end())
-        {
-            throw std::runtime_error("Cannot find \"" + key + "\" in joint_names.");
-        }
-        return static_cast<int>(std::distance(vec.begin(), it));
-    }
-
-    std::vector<std::vector<std::vector<std::array<float, 3>>>> update_keypoints(
-        const std::vector<std::vector<std::vector<std::array<float, 3>>>> &poses_2d,
-        const std::vector<std::string> &joint_names)
-    {
-        std::vector<std::vector<std::vector<std::array<float, 3>>>> new_views;
-        new_views.reserve(poses_2d.size());
-
-        for (const auto &view : poses_2d)
-        {
-            // "view" is a list of bodies => each body is Nx3
-            std::vector<std::vector<std::array<float, 3>>> new_bodies;
-            new_bodies.reserve(view.size());
-
-            for (const auto &body : view)
-            {
-                // 1) Copy first 17 keypoints
-                std::vector<std::array<float, 3>> new_body;
-                new_body.insert(new_body.end(), body.begin(), body.begin() + 17);
-
-                // 2) Optionally append extra keypoints
-                if (whole_body.at("foots"))
-                {
-                    new_body.insert(new_body.end(), body.begin() + 17, body.begin() + 23);
-                }
-                if (whole_body.at("face"))
-                {
-                    new_body.insert(new_body.end(), body.begin() + 23, body.begin() + 91);
-                }
-                if (whole_body.at("hands"))
-                {
-                    new_body.insert(new_body.end(), body.begin() + 91, body.end());
-                }
-
-                // 3) Compute mid_hip
-                int hlid = find_index(joint_names, "hip_left");
-                int hrid = find_index(joint_names, "hip_right");
-                float mid_hip_x = 0.5 * (new_body[hlid][0] + new_body[hrid][0]);
-                float mid_hip_y = 0.5 * (new_body[hlid][1] + new_body[hrid][1]);
-                float mid_hip_c = std::min(new_body[hlid][2], new_body[hrid][2]);
-                new_body.push_back({mid_hip_x, mid_hip_y, mid_hip_c});
-
-                // 4) Compute mid_shoulder
-                int slid = find_index(joint_names, "shoulder_left");
-                int srid = find_index(joint_names, "shoulder_right");
-                float mid_shoulder_x = 0.5 * (new_body[slid][0] + new_body[srid][0]);
-                float mid_shoulder_y = 0.5 * (new_body[slid][1] + new_body[srid][1]);
-                float mid_shoulder_c = std::min(new_body[slid][2], new_body[srid][2]);
-                new_body.push_back({mid_shoulder_x, mid_shoulder_y, mid_shoulder_c});
-
-                // 5) Compute head
-                int elid = find_index(joint_names, "ear_left");
-                int erid = find_index(joint_names, "ear_right");
-                float head_x = 0.5 * (new_body[elid][0] + new_body[erid][0]);
-                float head_y = 0.5 * (new_body[elid][1] + new_body[erid][1]);
-                float head_c = std::min(new_body[elid][2], new_body[erid][2]);
-                new_body.push_back({head_x, head_y, head_c});
-
-                // Add this updated body into new_bodies
-                new_bodies.push_back(new_body);
-            }
-
-            // Add all updated bodies for this view
-            new_views.push_back(new_bodies);
-        }
-
-        return new_views;
-    }
-}

scripts/test_skelda_dataset.py

@@ -8,8 +8,8 @@ import matplotlib
 import numpy as np
 import tqdm
 
-import test_triangulate
 import utils_2d_pose
+import utils_pipeline
 from skelda import evals
 
 sys.path.append("/RapidPoseTriangulation/swig/")
@@ -17,6 +17,12 @@ import rpt
 
 # ==================================================================================================
 
+whole_body = {
+    "foots": False,
+    "face": False,
+    "hands": False,
+}
+
 dataset_use = "human36m"
 # dataset_use = "panoptic"
 # dataset_use = "mvor"
@@ -175,7 +181,7 @@ datasets = {
     },
 }
 
-joint_names_2d = test_triangulate.joint_names_2d
+joint_names_2d = utils_pipeline.get_joint_names(whole_body)
 joint_names_3d = list(joint_names_2d)
 eval_joints = [
     "head",
@@ -197,7 +203,7 @@ if dataset_use == "human36m":
 if dataset_use == "panoptic":
     eval_joints[eval_joints.index("head")] = "nose"
 if dataset_use == "human36m_wb":
-    if any((test_triangulate.whole_body.values())):
+    if utils_pipeline.use_whole_body(whole_body):
         eval_joints = list(joint_names_2d)
     else:
         eval_joints[eval_joints.index("head")] = "nose"
@@ -323,9 +329,8 @@ def main():
     batch_poses = datasets[dataset_use].get("batch_poses", default_batch_poses)
 
     # Load 2D pose model
-    whole_body = test_triangulate.whole_body
+    if utils_pipeline.use_whole_body(whole_body):
-    if any((whole_body[k] for k in whole_body)):
+        kpt_model = utils_2d_pose.load_wb_model(min_bbox_score, min_bbox_area, batch_poses)
-        kpt_model = utils_2d_pose.load_wb_model()
     else:
         kpt_model = utils_2d_pose.load_model(min_bbox_score, min_bbox_area, batch_poses)
 
@@ -354,7 +359,7 @@ def main():
         try:
             for i in range(len(label["imgpaths"])):
                 imgpath = label["imgpaths"][i]
-                img = test_triangulate.load_image(imgpath)
+                img = utils_pipeline.load_image(imgpath)
         except cv2.error:
             print("One of the paths not found:", label["imgpaths"])
             continue
@@ -370,7 +375,7 @@ def main():
         try:
             for i in range(len(label["imgpaths"])):
                 imgpath = label["imgpaths"][i]
-                img = test_triangulate.load_image(imgpath)
+                img = utils_pipeline.load_image(imgpath)
                 images_2d.append(img)
         except cv2.error:
             print("One of the paths not found:", label["imgpaths"])
@@ -393,14 +398,14 @@ def main():
         # This also resulted in notably better MPJPE results in most cases, presumbly since the
         # demosaicing algorithm from OpenCV is better than the default one from the cameras
         for i in range(len(images_2d)):
-            images_2d[i] = test_triangulate.rgb2bayer(images_2d[i])
+            images_2d[i] = utils_pipeline.rgb2bayer(images_2d[i])
         time_imgs = time.time() - start
 
         start = time.time()
         for i in range(len(images_2d)):
-            images_2d[i] = test_triangulate.bayer2rgb(images_2d[i])
+            images_2d[i] = utils_pipeline.bayer2rgb(images_2d[i])
         poses_2d = utils_2d_pose.get_2d_pose(kpt_model, images_2d)
-        poses_2d = test_triangulate.update_keypoints(poses_2d, joint_names_2d)
+        poses_2d = utils_pipeline.update_keypoints(poses_2d, joint_names_2d, whole_body)
         time_2d = time.time() - start
 
         all_poses_2d.append(poses_2d)

scripts/test_triangulate.py

@@ -3,13 +3,12 @@ import json
 import os
 import sys
 import time
-from typing import List
 
-import cv2
 import matplotlib
 import numpy as np
 
 import utils_2d_pose
+import utils_pipeline
 from skelda import utils_pose, utils_view
 
 sys.path.append("/RapidPoseTriangulation/swig/")
@@ -25,175 +24,9 @@ whole_body = {
     "hands": False,
 }
 
-joint_names_2d = [
+joint_names_2d = utils_pipeline.get_joint_names(whole_body)
-    "nose",
-    "eye_left",
-    "eye_right",
-    "ear_left",
-    "ear_right",
-    "shoulder_left",
-    "shoulder_right",
-    "elbow_left",
-    "elbow_right",
-    "wrist_left",
-    "wrist_right",
-    "hip_left",
-    "hip_right",
-    "knee_left",
-    "knee_right",
-    "ankle_left",
-    "ankle_right",
-]
-if whole_body["foots"]:
-    joint_names_2d.extend(
-        [
-            "foot_toe_big_left",
-            "foot_toe_small_left",
-            "foot_heel_left",
-            "foot_toe_big_right",
-            "foot_toe_small_right",
-            "foot_heel_right",
-        ]
-    )
-if whole_body["face"]:
-    joint_names_2d.extend(
-        [
-            "face_jaw_right_1",
-            "face_jaw_right_2",
-            "face_jaw_right_3",
-            "face_jaw_right_4",
-            "face_jaw_right_5",
-            "face_jaw_right_6",
-            "face_jaw_right_7",
-            "face_jaw_right_8",
-            "face_jaw_middle",
-            "face_jaw_left_1",
-            "face_jaw_left_2",
-            "face_jaw_left_3",
-            "face_jaw_left_4",
-            "face_jaw_left_5",
-            "face_jaw_left_6",
-            "face_jaw_left_7",
-            "face_jaw_left_8",
-            "face_eyebrow_right_1",
-            "face_eyebrow_right_2",
-            "face_eyebrow_right_3",
-            "face_eyebrow_right_4",
-            "face_eyebrow_right_5",
-            "face_eyebrow_left_1",
-            "face_eyebrow_left_2",
-            "face_eyebrow_left_3",
-            "face_eyebrow_left_4",
-            "face_eyebrow_left_5",
-            "face_nose_1",
-            "face_nose_2",
-            "face_nose_3",
-            "face_nose_4",
-            "face_nose_5",
-            "face_nose_6",
-            "face_nose_7",
-            "face_nose_8",
-            "face_nose_9",
-            "face_eye_right_1",
-            "face_eye_right_2",
-            "face_eye_right_3",
-            "face_eye_right_4",
-            "face_eye_right_5",
-            "face_eye_right_6",
-            "face_eye_left_1",
-            "face_eye_left_2",
-            "face_eye_left_3",
-            "face_eye_left_4",
-            "face_eye_left_5",
-            "face_eye_left_6",
-            "face_mouth_1",
-            "face_mouth_2",
-            "face_mouth_3",
-            "face_mouth_4",
-            "face_mouth_5",
-            "face_mouth_6",
-            "face_mouth_7",
-            "face_mouth_8",
-            "face_mouth_9",
-            "face_mouth_10",
-            "face_mouth_11",
-            "face_mouth_12",
-            "face_mouth_13",
-            "face_mouth_14",
-            "face_mouth_15",
-            "face_mouth_16",
-            "face_mouth_17",
-            "face_mouth_18",
-            "face_mouth_19",
-            "face_mouth_20",
-        ]
-    )
-if whole_body["hands"]:
-    joint_names_2d.extend(
-        [
-            "hand_wrist_left",
-            "hand_finger_thumb_left_1",
-            "hand_finger_thumb_left_2",
-            "hand_finger_thumb_left_3",
-            "hand_finger_thumb_left_4",
-            "hand_finger_index_left_1",
-            "hand_finger_index_left_2",
-            "hand_finger_index_left_3",
-            "hand_finger_index_left_4",
-            "hand_finger_middle_left_1",
-            "hand_finger_middle_left_2",
-            "hand_finger_middle_left_3",
-            "hand_finger_middle_left_4",
-            "hand_finger_ring_left_1",
-            "hand_finger_ring_left_2",
-            "hand_finger_ring_left_3",
-            "hand_finger_ring_left_4",
-            "hand_finger_pinky_left_1",
-            "hand_finger_pinky_left_2",
-            "hand_finger_pinky_left_3",
-            "hand_finger_pinky_left_4",
-            "hand_wrist_right",
-            "hand_finger_thumb_right_1",
-            "hand_finger_thumb_right_2",
-            "hand_finger_thumb_right_3",
-            "hand_finger_thumb_right_4",
-            "hand_finger_index_right_1",
-            "hand_finger_index_right_2",
-            "hand_finger_index_right_3",
-            "hand_finger_index_right_4",
-            "hand_finger_middle_right_1",
-            "hand_finger_middle_right_2",
-            "hand_finger_middle_right_3",
-            "hand_finger_middle_right_4",
-            "hand_finger_ring_right_1",
-            "hand_finger_ring_right_2",
-            "hand_finger_ring_right_3",
-            "hand_finger_ring_right_4",
-            "hand_finger_pinky_right_1",
-            "hand_finger_pinky_right_2",
-            "hand_finger_pinky_right_3",
-            "hand_finger_pinky_right_4",
-        ]
-    )
-joint_names_2d.extend(
-    [
-        "hip_middle",
-        "shoulder_middle",
-        "head",
-    ]
-)
 joint_names_3d = list(joint_names_2d)
 
-main_limbs = [
-    ("shoulder_left", "elbow_left"),
-    ("elbow_left", "wrist_left"),
-    ("shoulder_right", "elbow_right"),
-    ("elbow_right", "wrist_right"),
-    ("hip_left", "knee_left"),
-    ("knee_left", "ankle_left"),
-    ("hip_right", "knee_right"),
-    ("knee_right", "ankle_right"),
-]
 
 # ==================================================================================================
 
@@ -217,85 +50,6 @@ def update_sample(sample, new_dir=""):
 # ==================================================================================================
 
 
-def load_image(path: str):
-    image = cv2.imread(path, 3)
-    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    image = np.asarray(image, dtype=np.uint8)
-    return image
-
-
-# ==================================================================================================
-
-
-def rgb2bayer(img):
-    bayer = np.zeros((img.shape[0], img.shape[1]), dtype=img.dtype)
-    bayer[0::2, 0::2] = img[0::2, 0::2, 0]
-    bayer[0::2, 1::2] = img[0::2, 1::2, 1]
-    bayer[1::2, 0::2] = img[1::2, 0::2, 1]
-    bayer[1::2, 1::2] = img[1::2, 1::2, 2]
-    return bayer
-
-
-def bayer2rgb(bayer):
-    img = cv2.cvtColor(bayer, cv2.COLOR_BayerBG2RGB)
-    return img
-
-
-# ==================================================================================================
-
-
-def update_keypoints(poses_2d: list, joint_names: List[str]) -> list:
-    new_views = []
-    for view in poses_2d:
-        new_bodies = []
-        for body in view:
-            body = body.tolist()
-
-            new_body = body[:17]
-            if whole_body["foots"]:
-                new_body.extend(body[17:23])
-            if whole_body["face"]:
-                new_body.extend(body[23:91])
-            if whole_body["hands"]:
-                new_body.extend(body[91:])
-            body = new_body
-
-            hlid = joint_names.index("hip_left")
-            hrid = joint_names.index("hip_right")
-            mid_hip = [
-                float(((body[hlid][0] + body[hrid][0]) / 2.0)),
-                float(((body[hlid][1] + body[hrid][1]) / 2.0)),
-                min(body[hlid][2], body[hrid][2]),
-            ]
-            body.append(mid_hip)
-
-            slid = joint_names.index("shoulder_left")
-            srid = joint_names.index("shoulder_right")
-            mid_shoulder = [
-                float(((body[slid][0] + body[srid][0]) / 2.0)),
-                float(((body[slid][1] + body[srid][1]) / 2.0)),
-                min(body[slid][2], body[srid][2]),
-            ]
-            body.append(mid_shoulder)
-
-            elid = joint_names.index("ear_left")
-            erid = joint_names.index("ear_right")
-            head = [
-                float(((body[elid][0] + body[erid][0]) / 2.0)),
-                float(((body[elid][1] + body[erid][1]) / 2.0)),
-                min(body[elid][2], body[erid][2]),
-            ]
-            body.append(head)
-
-            new_bodies.append(body)
-        new_views.append(new_bodies)
-
-    return new_views
-
-
-# ==================================================================================================
-
-
 def main():
     if any((whole_body[k] for k in whole_body)):
         kpt_model = utils_2d_pose.load_wb_model()
@@ -330,15 +84,15 @@ def main():
         images_2d = []
         for i in range(len(sample["cameras_color"])):
             imgpath = sample["imgpaths_color"][i]
-            img = load_image(imgpath)
+            img = utils_pipeline.load_image(imgpath)
-            img = rgb2bayer(img)
+            img = utils_pipeline.rgb2bayer(img)
-            img = bayer2rgb(img)
+            img = utils_pipeline.bayer2rgb(img)
             images_2d.append(img)
 
         # Get 2D poses
         stime = time.time()
         poses_2d = utils_2d_pose.get_2d_pose(kpt_model, images_2d)
-        poses_2d = update_keypoints(poses_2d, joint_names_2d)
+        poses_2d = utils_pipeline.update_keypoints(poses_2d, joint_names_2d, whole_body)
         print("2D time:", time.time() - stime)
         # print([np.array(p).round(6).tolist() for p in poses_2d])
 

scripts/utils_pipeline.hpp

@@ -0,0 +1,316 @@
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include <opencv2/opencv.hpp>
+
+// =================================================================================================
+
+namespace utils_pipeline
+{
+    bool use_whole_body(const std::map<std::string, bool> &whole_body)
+    {
+        for (const auto &pair : whole_body)
+        {
+            if (pair.second)
+            {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    // =============================================================================================
+
+    std::vector<std::string> get_joint_names(const std::map<std::string, bool> &whole_body)
+    {
+        std::vector<std::string> joint_names_2d = {
+            "nose",
+            "eye_left",
+            "eye_right",
+            "ear_left",
+            "ear_right",
+            "shoulder_left",
+            "shoulder_right",
+            "elbow_left",
+            "elbow_right",
+            "wrist_left",
+            "wrist_right",
+            "hip_left",
+            "hip_right",
+            "knee_left",
+            "knee_right",
+            "ankle_left",
+            "ankle_right",
+        };
+
+        if (whole_body.at("foots"))
+        {
+            joint_names_2d.insert(
+                joint_names_2d.end(),
+                {
+                    "foot_toe_big_left",
+                    "foot_toe_small_left",
+                    "foot_heel_left",
+                    "foot_toe_big_right",
+                    "foot_toe_small_right",
+                    "foot_heel_right",
+                });
+        }
+        if (whole_body.at("face"))
+        {
+            joint_names_2d.insert(
+                joint_names_2d.end(),
+                {
+                    "face_jaw_right_1",
+                    "face_jaw_right_2",
+                    "face_jaw_right_3",
+                    "face_jaw_right_4",
+                    "face_jaw_right_5",
+                    "face_jaw_right_6",
+                    "face_jaw_right_7",
+                    "face_jaw_right_8",
+                    "face_jaw_middle",
+                    "face_jaw_left_1",
+                    "face_jaw_left_2",
+                    "face_jaw_left_3",
+                    "face_jaw_left_4",
+                    "face_jaw_left_5",
+                    "face_jaw_left_6",
+                    "face_jaw_left_7",
+                    "face_jaw_left_8",
+                    "face_eyebrow_right_1",
+                    "face_eyebrow_right_2",
+                    "face_eyebrow_right_3",
+                    "face_eyebrow_right_4",
+                    "face_eyebrow_right_5",
+                    "face_eyebrow_left_1",
+                    "face_eyebrow_left_2",
+                    "face_eyebrow_left_3",
+                    "face_eyebrow_left_4",
+                    "face_eyebrow_left_5",
+                    "face_nose_1",
+                    "face_nose_2",
+                    "face_nose_3",
+                    "face_nose_4",
+                    "face_nose_5",
+                    "face_nose_6",
+                    "face_nose_7",
+                    "face_nose_8",
+                    "face_nose_9",
+                    "face_eye_right_1",
+                    "face_eye_right_2",
+                    "face_eye_right_3",
+                    "face_eye_right_4",
+                    "face_eye_right_5",
+                    "face_eye_right_6",
+                    "face_eye_left_1",
+                    "face_eye_left_2",
+                    "face_eye_left_3",
+                    "face_eye_left_4",
+                    "face_eye_left_5",
+                    "face_eye_left_6",
+                    "face_mouth_1",
+                    "face_mouth_2",
+                    "face_mouth_3",
+                    "face_mouth_4",
+                    "face_mouth_5",
+                    "face_mouth_6",
+                    "face_mouth_7",
+                    "face_mouth_8",
+                    "face_mouth_9",
+                    "face_mouth_10",
+                    "face_mouth_11",
+                    "face_mouth_12",
+                    "face_mouth_13",
+                    "face_mouth_14",
+                    "face_mouth_15",
+                    "face_mouth_16",
+                    "face_mouth_17",
+                    "face_mouth_18",
+                    "face_mouth_19",
+                    "face_mouth_20",
+                });
+        }
+        if (whole_body.at("hands"))
+        {
+            joint_names_2d.insert(
+                joint_names_2d.end(),
+                {
+                    "hand_wrist_left",
+                    "hand_finger_thumb_left_1",
+                    "hand_finger_thumb_left_2",
+                    "hand_finger_thumb_left_3",
+                    "hand_finger_thumb_left_4",
+                    "hand_finger_index_left_1",
+                    "hand_finger_index_left_2",
+                    "hand_finger_index_left_3",
+                    "hand_finger_index_left_4",
+                    "hand_finger_middle_left_1",
+                    "hand_finger_middle_left_2",
+                    "hand_finger_middle_left_3",
+                    "hand_finger_middle_left_4",
+                    "hand_finger_ring_left_1",
+                    "hand_finger_ring_left_2",
+                    "hand_finger_ring_left_3",
+                    "hand_finger_ring_left_4",
+                    "hand_finger_pinky_left_1",
+                    "hand_finger_pinky_left_2",
+                    "hand_finger_pinky_left_3",
+                    "hand_finger_pinky_left_4",
+                    "hand_wrist_right",
+                    "hand_finger_thumb_right_1",
+                    "hand_finger_thumb_right_2",
+                    "hand_finger_thumb_right_3",
+                    "hand_finger_thumb_right_4",
+                    "hand_finger_index_right_1",
+                    "hand_finger_index_right_2",
+                    "hand_finger_index_right_3",
+                    "hand_finger_index_right_4",
+                    "hand_finger_middle_right_1",
+                    "hand_finger_middle_right_2",
+                    "hand_finger_middle_right_3",
+                    "hand_finger_middle_right_4",
+                    "hand_finger_ring_right_1",
+                    "hand_finger_ring_right_2",
+                    "hand_finger_ring_right_3",
+                    "hand_finger_ring_right_4",
+                    "hand_finger_pinky_right_1",
+                    "hand_finger_pinky_right_2",
+                    "hand_finger_pinky_right_3",
+                    "hand_finger_pinky_right_4",
+                });
+        }
+
+        joint_names_2d.insert(
+            joint_names_2d.end(),
+            {
+                "hip_middle",
+                "shoulder_middle",
+                "head",
+            });
+
+        return joint_names_2d;
+    }
+
+    // =============================================================================================
+
+    cv::Mat bayer2rgb(const cv::Mat &bayer)
+    {
+        cv::Mat rgb;
+        cv::cvtColor(bayer, rgb, cv::COLOR_BayerBG2RGB);
+        return rgb;
+    }
+
+    cv::Mat rgb2bayer(const cv::Mat &img)
+    {
+        CV_Assert(img.type() == CV_8UC3);
+        cv::Mat bayer(img.rows, img.cols, CV_8UC1);
+
+        for (int r = 0; r < img.rows; ++r)
+        {
+            const uchar *imgData = img.ptr<uchar>(r);
+            uchar *bayerRowPtr = bayer.ptr<uchar>(r);
+
+            for (int c = 0; c < img.cols; ++c)
+            {
+                int pixelIndex = 3 * c;
+
+                // Use faster bit operation instead of modulo+if
+                // Even row, even col => R = 0
+                // Even row, odd col => G = 1
+                // Odd row, even col => G = 1
+                // Odd row, odd col => B = 2
+                int row_mod = r & 1;
+                int col_mod = c & 1;
+                int component = row_mod + col_mod;
+
+                bayerRowPtr[c] = imgData[pixelIndex + component];
+            }
+        }
+
+        return bayer;
+    }
+
+    // =============================================================================================
+
+    inline int find_index(const std::vector<std::string> &vec, const std::string &key)
+    {
+        auto it = std::find(vec.begin(), vec.end(), key);
+        if (it == vec.end())
+        {
+            throw std::runtime_error("Cannot find \"" + key + "\" in joint_names.");
+        }
+        return static_cast<int>(std::distance(vec.begin(), it));
+    }
+
+    std::vector<std::vector<std::vector<std::array<float, 3>>>> update_keypoints(
+        const std::vector<std::vector<std::vector<std::array<float, 3>>>> &poses_2d,
+        const std::vector<std::string> &joint_names,
+        const std::map<std::string, bool> &whole_body)
+    {
+        std::vector<std::vector<std::vector<std::array<float, 3>>>> new_views;
+        new_views.reserve(poses_2d.size());
+
+        for (const auto &view : poses_2d)
+        {
+            // "view" is a list of bodies => each body is Nx3
+            std::vector<std::vector<std::array<float, 3>>> new_bodies;
+            new_bodies.reserve(view.size());
+
+            for (const auto &body : view)
+            {
+                // 1) Copy first 17 keypoints
+                std::vector<std::array<float, 3>> new_body;
+                new_body.insert(new_body.end(), body.begin(), body.begin() + 17);
+
+                // 2) Optionally append extra keypoints
+                if (whole_body.at("foots"))
+                {
+                    new_body.insert(new_body.end(), body.begin() + 17, body.begin() + 23);
+                }
+                if (whole_body.at("face"))
+                {
+                    new_body.insert(new_body.end(), body.begin() + 23, body.begin() + 91);
+                }
+                if (whole_body.at("hands"))
+                {
+                    new_body.insert(new_body.end(), body.begin() + 91, body.end());
+                }
+
+                // 3) Compute mid_hip
+                int hlid = find_index(joint_names, "hip_left");
+                int hrid = find_index(joint_names, "hip_right");
+                float mid_hip_x = 0.5 * (new_body[hlid][0] + new_body[hrid][0]);
+                float mid_hip_y = 0.5 * (new_body[hlid][1] + new_body[hrid][1]);
+                float mid_hip_c = std::min(new_body[hlid][2], new_body[hrid][2]);
+                new_body.push_back({mid_hip_x, mid_hip_y, mid_hip_c});
+
+                // 4) Compute mid_shoulder
+                int slid = find_index(joint_names, "shoulder_left");
+                int srid = find_index(joint_names, "shoulder_right");
+                float mid_shoulder_x = 0.5 * (new_body[slid][0] + new_body[srid][0]);
+                float mid_shoulder_y = 0.5 * (new_body[slid][1] + new_body[srid][1]);
+                float mid_shoulder_c = std::min(new_body[slid][2], new_body[srid][2]);
+                new_body.push_back({mid_shoulder_x, mid_shoulder_y, mid_shoulder_c});
+
+                // 5) Compute head
+                int elid = find_index(joint_names, "ear_left");
+                int erid = find_index(joint_names, "ear_right");
+                float head_x = 0.5 * (new_body[elid][0] + new_body[erid][0]);
+                float head_y = 0.5 * (new_body[elid][1] + new_body[erid][1]);
+                float head_c = std::min(new_body[elid][2], new_body[erid][2]);
+                new_body.push_back({head_x, head_y, head_c});
+
+                // Add this updated body into new_bodies
+                new_bodies.push_back(new_body);
+            }
+
+            // Add all updated bodies for this view
+            new_views.push_back(new_bodies);
+        }
+
+        return new_views;
+    }
+}

scripts/utils_pipeline.py

@@ -0,0 +1,255 @@
+from typing import List
+
+import cv2
+import numpy as np
+
+# ==================================================================================================
+
+def use_whole_body(whole_body: dict) -> bool:
+    return any((whole_body[k] for k in whole_body))
+
+# ==================================================================================================
+
+
+def get_joint_names(whole_body: dict):
+
+    joint_names_2d = [
+        "nose",
+        "eye_left",
+        "eye_right",
+        "ear_left",
+        "ear_right",
+        "shoulder_left",
+        "shoulder_right",
+        "elbow_left",
+        "elbow_right",
+        "wrist_left",
+        "wrist_right",
+        "hip_left",
+        "hip_right",
+        "knee_left",
+        "knee_right",
+        "ankle_left",
+        "ankle_right",
+    ]
+
+    if whole_body["foots"]:
+        joint_names_2d.extend(
+            [
+                "foot_toe_big_left",
+                "foot_toe_small_left",
+                "foot_heel_left",
+                "foot_toe_big_right",
+                "foot_toe_small_right",
+                "foot_heel_right",
+            ]
+        )
+    if whole_body["face"]:
+        joint_names_2d.extend(
+            [
+                "face_jaw_right_1",
+                "face_jaw_right_2",
+                "face_jaw_right_3",
+                "face_jaw_right_4",
+                "face_jaw_right_5",
+                "face_jaw_right_6",
+                "face_jaw_right_7",
+                "face_jaw_right_8",
+                "face_jaw_middle",
+                "face_jaw_left_1",
+                "face_jaw_left_2",
+                "face_jaw_left_3",
+                "face_jaw_left_4",
+                "face_jaw_left_5",
+                "face_jaw_left_6",
+                "face_jaw_left_7",
+                "face_jaw_left_8",
+                "face_eyebrow_right_1",
+                "face_eyebrow_right_2",
+                "face_eyebrow_right_3",
+                "face_eyebrow_right_4",
+                "face_eyebrow_right_5",
+                "face_eyebrow_left_1",
+                "face_eyebrow_left_2",
+                "face_eyebrow_left_3",
+                "face_eyebrow_left_4",
+                "face_eyebrow_left_5",
+                "face_nose_1",
+                "face_nose_2",
+                "face_nose_3",
+                "face_nose_4",
+                "face_nose_5",
+                "face_nose_6",
+                "face_nose_7",
+                "face_nose_8",
+                "face_nose_9",
+                "face_eye_right_1",
+                "face_eye_right_2",
+                "face_eye_right_3",
+                "face_eye_right_4",
+                "face_eye_right_5",
+                "face_eye_right_6",
+                "face_eye_left_1",
+                "face_eye_left_2",
+                "face_eye_left_3",
+                "face_eye_left_4",
+                "face_eye_left_5",
+                "face_eye_left_6",
+                "face_mouth_1",
+                "face_mouth_2",
+                "face_mouth_3",
+                "face_mouth_4",
+                "face_mouth_5",
+                "face_mouth_6",
+                "face_mouth_7",
+                "face_mouth_8",
+                "face_mouth_9",
+                "face_mouth_10",
+                "face_mouth_11",
+                "face_mouth_12",
+                "face_mouth_13",
+                "face_mouth_14",
+                "face_mouth_15",
+                "face_mouth_16",
+                "face_mouth_17",
+                "face_mouth_18",
+                "face_mouth_19",
+                "face_mouth_20",
+            ]
+        )
+    if whole_body["hands"]:
+        joint_names_2d.extend(
+            [
+                "hand_wrist_left",
+                "hand_finger_thumb_left_1",
+                "hand_finger_thumb_left_2",
+                "hand_finger_thumb_left_3",
+                "hand_finger_thumb_left_4",
+                "hand_finger_index_left_1",
+                "hand_finger_index_left_2",
+                "hand_finger_index_left_3",
+                "hand_finger_index_left_4",
+                "hand_finger_middle_left_1",
+                "hand_finger_middle_left_2",
+                "hand_finger_middle_left_3",
+                "hand_finger_middle_left_4",
+                "hand_finger_ring_left_1",
+                "hand_finger_ring_left_2",
+                "hand_finger_ring_left_3",
+                "hand_finger_ring_left_4",
+                "hand_finger_pinky_left_1",
+                "hand_finger_pinky_left_2",
+                "hand_finger_pinky_left_3",
+                "hand_finger_pinky_left_4",
+                "hand_wrist_right",
+                "hand_finger_thumb_right_1",
+                "hand_finger_thumb_right_2",
+                "hand_finger_thumb_right_3",
+                "hand_finger_thumb_right_4",
+                "hand_finger_index_right_1",
+                "hand_finger_index_right_2",
+                "hand_finger_index_right_3",
+                "hand_finger_index_right_4",
+                "hand_finger_middle_right_1",
+                "hand_finger_middle_right_2",
+                "hand_finger_middle_right_3",
+                "hand_finger_middle_right_4",
+                "hand_finger_ring_right_1",
+                "hand_finger_ring_right_2",
+                "hand_finger_ring_right_3",
+                "hand_finger_ring_right_4",
+                "hand_finger_pinky_right_1",
+                "hand_finger_pinky_right_2",
+                "hand_finger_pinky_right_3",
+                "hand_finger_pinky_right_4",
+            ]
+        )
+
+    joint_names_2d.extend(
+        [
+            "hip_middle",
+            "shoulder_middle",
+            "head",
+        ]
+    )
+
+    return joint_names_2d
+
+
+# ==================================================================================================
+
+
+def load_image(path: str):
+    image = cv2.imread(path, 3)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    image = np.asarray(image, dtype=np.uint8)
+    return image
+
+
+# ==================================================================================================
+
+
+def rgb2bayer(img):
+    bayer = np.zeros((img.shape[0], img.shape[1]), dtype=img.dtype)
+    bayer[0::2, 0::2] = img[0::2, 0::2, 0]
+    bayer[0::2, 1::2] = img[0::2, 1::2, 1]
+    bayer[1::2, 0::2] = img[1::2, 0::2, 1]
+    bayer[1::2, 1::2] = img[1::2, 1::2, 2]
+    return bayer
+
+
+def bayer2rgb(bayer):
+    img = cv2.cvtColor(bayer, cv2.COLOR_BayerBG2RGB)
+    return img
+
+
+# ==================================================================================================
+
+
+def update_keypoints(poses_2d: list, joint_names: List[str], whole_body: dict) -> list:
+    new_views = []
+    for view in poses_2d:
+        new_bodies = []
+        for body in view:
+            body = body.tolist()
+
+            new_body = body[:17]
+            if whole_body["foots"]:
+                new_body.extend(body[17:23])
+            if whole_body["face"]:
+                new_body.extend(body[23:91])
+            if whole_body["hands"]:
+                new_body.extend(body[91:])
+            body = new_body
+
+            hlid = joint_names.index("hip_left")
+            hrid = joint_names.index("hip_right")
+            mid_hip = [
+                float(((body[hlid][0] + body[hrid][0]) / 2.0)),
+                float(((body[hlid][1] + body[hrid][1]) / 2.0)),
+                min(body[hlid][2], body[hrid][2]),
+            ]
+            body.append(mid_hip)
+
+            slid = joint_names.index("shoulder_left")
+            srid = joint_names.index("shoulder_right")
+            mid_shoulder = [
+                float(((body[slid][0] + body[srid][0]) / 2.0)),
+                float(((body[slid][1] + body[srid][1]) / 2.0)),
+                min(body[slid][2], body[srid][2]),
+            ]
+            body.append(mid_shoulder)
+
+            elid = joint_names.index("ear_left")
+            erid = joint_names.index("ear_right")
+            head = [
+                float(((body[elid][0] + body[erid][0]) / 2.0)),
+                float(((body[elid][1] + body[erid][1]) / 2.0)),
+                min(body[elid][2], body[erid][2]),
+            ]
+            body.append(head)
+
+            new_bodies.append(body)
+        new_views.append(new_bodies)
+
+    return new_views

tests/README.md

@@ -1,28 +1,17 @@
-# ROS-Wrapper
+# Tests
 
-Run pose estimator with ros topics as inputs and publish detected poses.
+Various module tests
 
-<br>
+### Triangulator
 
-- Build container:
   ```bash
-  docker build --progress=plain -t rapidposetriangulation_ros -f ros/dockerfile .
+  cd /RapidPoseTriangulation/tests/ && python3 test_interface.py && cd ..
   ```
 
-- Update or remove the `ROS_DOMAIN_ID` in `docker-compose.yml` to fit your environment
+### Onnx C++ Interface
-
-- Run and test:
-  ```bash
-  xhost +; docker compose -f ros/docker-compose.yml up
-
-  docker exec -it ros-test_node-1 bash
-  export ROS_DOMAIN_ID=18
-  ```
-
-### Debugging
 
 ```bash
-cd /RapidPoseTriangulation/ros/rpt2D_wrapper_cpp/tests/
+cd /RapidPoseTriangulation/tests/
 
 g++ -std=c++17 -O3 -march=native -Wall -DCOMPILE_EXAMPLE_MAIN \
     $(pkg-config --cflags opencv4) \
@@ -30,7 +19,7 @@ g++ -std=c++17 -O3 -march=native -Wall -DCOMPILE_EXAMPLE_MAIN \
     -I /onnxruntime/include/onnxruntime/core/session \
     -I /onnxruntime/include/onnxruntime/core/providers/tensorrt \
     -L /onnxruntime/build/Linux/Release \
-    my_app.cpp \
+    test_utils2d.cpp \
     -o my_app \
     -Wl,--start-group \
       -lonnxruntime_providers_tensorrt \

tests/test_utils2d.cpp

@@ -5,7 +5,7 @@
 
 #include <opencv2/opencv.hpp>
 
-#include "../src/utils_2d_pose.hpp"
+#include "/RapidPoseTriangulation/scripts/utils_2d_pose.hpp"
 
 // =================================================================================================
 // =================================================================================================