Restructuring some code.
This commit is contained in:
Daniel
@ -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])
|
||||
|
||||
|
||||
Reference in New Issue
Block a user