Restructuring some code.
This commit is contained in:
@ -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 any((whole_body[k] for k in whole_body)):
|
||||
kpt_model = utils_2d_pose.load_wb_model()
|
||||
if utils_pipeline.use_whole_body(whole_body):
|
||||
kpt_model = utils_2d_pose.load_wb_model(min_bbox_score, min_bbox_area, batch_poses)
|
||||
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)
|
||||
|
||||
@ -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 = [
|
||||
"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_2d = utils_pipeline.get_joint_names(whole_body)
|
||||
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 = rgb2bayer(img)
|
||||
img = bayer2rgb(img)
|
||||
img = utils_pipeline.load_image(imgpath)
|
||||
img = utils_pipeline.rgb2bayer(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])
|
||||
|
||||
|
||||
1068
scripts/utils_2d_pose.hpp
Normal file
1068
scripts/utils_2d_pose.hpp
Normal file
File diff suppressed because it is too large
Load Diff
316
scripts/utils_pipeline.hpp
Normal file
316
scripts/utils_pipeline.hpp
Normal file
@ -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;
|
||||
}
|
||||
}
|
||||
255
scripts/utils_pipeline.py
Normal file
255
scripts/utils_pipeline.py
Normal file
@ -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
|
||||
Reference in New Issue
Block a user