145 lines
4.2 KiB
Python
145 lines
4.2 KiB
Python
import copy
|
|
import json
|
|
import os
|
|
|
|
import numpy as np
|
|
|
|
import utils_pipeline
|
|
from skelda import utils_pose, utils_view
|
|
from skelda.writers import json_writer
|
|
|
|
# ==================================================================================================
|
|
|
|
filepath = os.path.dirname(os.path.realpath(__file__)) + "/"
|
|
test_img_dir = filepath + "../data/"
|
|
|
|
whole_body = {
|
|
"foots": False,
|
|
"face": False,
|
|
"hands": False,
|
|
}
|
|
config = {
|
|
"min_match_score": 0.94,
|
|
"min_group_size": 1,
|
|
"min_bbox_score": 0.3,
|
|
"min_bbox_area": 0.1 * 0.1,
|
|
"batch_poses": True,
|
|
"whole_body": whole_body,
|
|
"take_interval": 1,
|
|
"fps": -1,
|
|
"max_movement_speed": 0,
|
|
"max_track_distance": 0,
|
|
}
|
|
|
|
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
|
|
if not "scene" in sample:
|
|
sample["scene"] = "default"
|
|
if not "id" in sample:
|
|
sample["id"] = "0"
|
|
if not "index" in sample:
|
|
sample["index"] = 0
|
|
for cam in sample["cameras"]:
|
|
if not "type" in cam:
|
|
cam["type"] = "pinhole"
|
|
|
|
return sample
|
|
|
|
|
|
# ==================================================================================================
|
|
|
|
|
|
def main():
|
|
|
|
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)
|
|
|
|
if len(sample["imgpaths"]) == 1:
|
|
# At least two images are required
|
|
continue
|
|
|
|
# Save dataset
|
|
labels = [sample]
|
|
tmp_export_dir = "/tmp/rpt/"
|
|
for label in labels:
|
|
if "splits" in label:
|
|
label.pop("splits")
|
|
json_writer.save_dataset(labels, tmp_export_dir)
|
|
|
|
# Save config
|
|
config_path = tmp_export_dir + "config.json"
|
|
utils_pipeline.save_json(config, config_path)
|
|
|
|
# Call the CPP binary
|
|
os.system("/RapidPoseTriangulation/scripts/test_skelda_dataset.bin")
|
|
|
|
# Load the results
|
|
print("Loading exports ...")
|
|
res_path = tmp_export_dir + "results.json"
|
|
results = utils_pipeline.load_json(res_path)
|
|
poses_3d = results["all_poses_3d"][0]
|
|
poses_2d = results["all_poses_2d"][0]
|
|
joint_names_3d = results["joint_names_3d"]
|
|
|
|
# Visualize the 2D results
|
|
fig1 = utils_view.draw_many_images(
|
|
sample["imgpaths"], [], [], poses_2d, joint_names_2d, "2D detections"
|
|
)
|
|
fig1.savefig(os.path.join(dirpath, "2d-k.png"), dpi=fig1.dpi)
|
|
|
|
# Visualize the 3D results
|
|
print("Detected 3D poses:")
|
|
poses_3d = np.array(poses_3d)
|
|
print(poses_3d.round(3))
|
|
if len(poses_3d) == 0:
|
|
utils_view.show_plots()
|
|
continue
|
|
camparams = sample["cameras"]
|
|
roomparams = {
|
|
"room_size": sample["room_size"],
|
|
"room_center": sample["room_center"],
|
|
}
|
|
poses_2d_proj = []
|
|
for cam in camparams:
|
|
poses_2d_cam, _ = utils_pose.project_poses(poses_3d, cam)
|
|
poses_2d_proj.append(poses_2d_cam)
|
|
fig2 = utils_view.draw_poses3d(poses_3d, joint_names_3d, roomparams, camparams)
|
|
fig3 = utils_view.draw_many_images(
|
|
sample["imgpaths"], [], [], poses_2d_proj, 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()
|