Restructuring some code.

scripts/test_skelda_dataset.py

@@ -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)