317 lines
11 KiB
C++
317 lines
11 KiB
C++
#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;
|
|
}
|
|
}
|