149 lines
4.6 KiB
Python
149 lines
4.6 KiB
Python
import copy
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
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/")
|
|
import rpt
|
|
|
|
# ==================================================================================================
|
|
|
|
filepath = os.path.dirname(os.path.realpath(__file__)) + "/"
|
|
test_img_dir = filepath + "../data/"
|
|
whole_body = {
|
|
"foots": False,
|
|
"face": False,
|
|
"hands": False,
|
|
}
|
|
|
|
joint_names_2d = utils_pipeline.get_joint_names(whole_body)
|
|
joint_names_3d = list(joint_names_2d)
|
|
|
|
|
|
# ==================================================================================================
|
|
|
|
|
|
def update_sample(sample, new_dir=""):
|
|
sample = copy.deepcopy(sample)
|
|
|
|
# Rename image paths
|
|
sample["imgpaths"] = [
|
|
os.path.join(new_dir, os.path.basename(v)) for v in sample["imgpaths"]
|
|
]
|
|
|
|
# Add placeholders for missing keys
|
|
sample["cameras_color"] = sample["cameras"]
|
|
sample["imgpaths_color"] = sample["imgpaths"]
|
|
sample["cameras_depth"] = []
|
|
|
|
return sample
|
|
|
|
|
|
# ==================================================================================================
|
|
|
|
|
|
def main():
|
|
if any((whole_body[k] for k in whole_body)):
|
|
kpt_model = utils_2d_pose.load_wb_model()
|
|
else:
|
|
kpt_model = utils_2d_pose.load_model(min_bbox_score=0.3)
|
|
|
|
# Manually set matplotlib backend
|
|
matplotlib.use("TkAgg")
|
|
|
|
for dirname in sorted(os.listdir(test_img_dir)):
|
|
dirpath = os.path.join(test_img_dir, dirname)
|
|
|
|
if not os.path.isdir(dirpath):
|
|
continue
|
|
|
|
if (dirname[0] not in ["p", "h", "e", "q"]) or len(dirname) != 2:
|
|
continue
|
|
|
|
# Load sample infos
|
|
print("\n" + dirpath)
|
|
with open(os.path.join(dirpath, "sample.json"), "r", encoding="utf-8") as file:
|
|
sample = json.load(file)
|
|
sample = update_sample(sample, dirpath)
|
|
|
|
camparams = sample["cameras_color"]
|
|
roomparams = {
|
|
"room_size": sample["room_size"],
|
|
"room_center": sample["room_center"],
|
|
}
|
|
|
|
# Load color images
|
|
images_2d = []
|
|
for i in range(len(sample["cameras_color"])):
|
|
imgpath = sample["imgpaths_color"][i]
|
|
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 = 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])
|
|
|
|
fig1 = utils_view.draw_many_images(
|
|
sample["imgpaths_color"], [], [], poses_2d, joint_names_2d, "2D detections"
|
|
)
|
|
fig1.savefig(os.path.join(dirpath, "2d-k.png"), dpi=fig1.dpi)
|
|
# draw_utils.utils_view.show_plots()
|
|
|
|
if len(images_2d) == 1:
|
|
utils_view.show_plots()
|
|
continue
|
|
|
|
# Get 3D poses
|
|
if sum(np.sum(p) for p in poses_2d) == 0:
|
|
poses3D = np.zeros([1, len(joint_names_3d), 4])
|
|
poses2D = np.zeros([len(images_2d), 1, len(joint_names_3d), 3])
|
|
else:
|
|
cameras = rpt.convert_cameras(camparams)
|
|
roomp = [roomparams["room_size"], roomparams["room_center"]]
|
|
triangulator = rpt.Triangulator(min_match_score=0.94)
|
|
|
|
stime = time.time()
|
|
poses_3d = triangulator.triangulate_poses(
|
|
poses_2d, cameras, roomp, joint_names_2d
|
|
)
|
|
poses3D = np.array(poses_3d)
|
|
if len(poses3D) == 0:
|
|
poses3D = np.zeros([1, len(joint_names_3d), 4])
|
|
print("3D time:", time.time() - stime)
|
|
|
|
poses2D = []
|
|
for cam in camparams:
|
|
poses_2d, _ = utils_pose.project_poses(poses3D, cam)
|
|
poses2D.append(poses_2d)
|
|
|
|
print(poses3D)
|
|
# print(poses2D)
|
|
# print(poses3D.round(3).tolist())
|
|
|
|
fig2 = utils_view.draw_poses3d(poses3D, joint_names_3d, roomparams, camparams)
|
|
fig3 = utils_view.draw_many_images(
|
|
sample["imgpaths_color"], [], [], poses2D, joint_names_3d, "2D projections"
|
|
)
|
|
fig2.savefig(os.path.join(dirpath, "3d-p.png"), dpi=fig2.dpi)
|
|
fig3.savefig(os.path.join(dirpath, "2d-p.png"), dpi=fig3.dpi)
|
|
utils_view.show_plots()
|
|
|
|
|
|
# ==================================================================================================
|
|
|
|
if __name__ == "__main__":
|
|
main()
|