Some mixed updates.

scripts/test_skelda_dataset.py

@@ -1,6 +1,8 @@
 import json
 import os
 
+import numpy as np
+
 import utils_pipeline
 from skelda import evals
 from skelda.writers import json_writer
@@ -249,6 +251,10 @@ if dataset_use == "human36m_wb":
 # output_dir = "/RapidPoseTriangulation/data/testoutput/"
 output_dir = ""
 
+# pred_export_path = f"/datasets/predictions/{dataset_use}/RapidPoseTriangulation.json"
+pred_export_path = ""
+
+
 # ==================================================================================================
 
 
@@ -281,10 +287,6 @@ def load_labels(dataset: dict):
         labels = [lb for lb in labels if lb["subject"] == "S9"]
         labels = [lb for i, lb in enumerate(labels) if i % 4000 < 150]
 
-        for label in labels:
-            label.pop("action")
-            label.pop("frame")
-
     elif "mvor" in dataset:
         labels = utils_pipeline.load_json(dataset["mvor"]["path"])
 
@@ -365,6 +367,7 @@ def main():
     )
 
     # Print a dataset sample for debugging
+    print("Amount of samples:", len(labels))
     print(labels[0])
 
     # Save dataset
@@ -417,6 +420,59 @@ def main():
     all_ids = results["all_ids"]
     joint_names_3d = results["joint_names_3d"]
 
+    # # Visualize labels and predictions
+    # from skelda import utils_view
+    # for i in range(0, len(labels), 1):
+    #     posesL = []
+    #     posesR = []
+    #     jnames = []
+    #     for j in labels[i]["joints"]:
+    #         if "->" in j:
+    #             jnames.append(j.split("->")[-1])
+    #         else:
+    #             jnames.append(j)
+    #     for j in range(len(labels[i]["bodies3D"])):
+    #         pose = []
+    #         for k in range(len(eval_joints)):
+    #             n = eval_joints[k]
+    #             pose.append(labels[i]["bodies3D"][j][jnames.index(n)])
+    #         posesL.append(pose)
+    #     for j in range(len(all_poses_3d[i])):
+    #         pose = []
+    #         for k in range(len(eval_joints)):
+    #             n = eval_joints[k]
+    #             pose.append(all_poses_3d[i][j][joint_names_3d.index(n)])
+    #         posesR.append(pose)
+    #     poses_3d = posesL + posesR
+    #     sample = labels[i]
+    #     sample["bodies3D"] = np.array(poses_3d).round(3).tolist()
+    #     sample["joints"] = eval_joints
+    #     sample["num_persons"] = len(poses_3d)
+    #     print(sample)
+    #     utils_view.draw_sample_3d(sample)
+    #     utils_view.draw_many_images(
+    #         sample["imgpaths"],
+    #         sample["cameras"],
+    #         [],
+    #         all_poses_2d[i],
+    #         joint_names_3d,
+    #         "2D detections",
+    #     )
+    #     utils_view.show_plots()
+
+    if pred_export_path != "":
+        # Export predictions
+        print("\nExporting predictions ...")
+        all_poses_3d = [np.array(poses).round(3).tolist() for poses in all_poses_3d]
+        data = {
+            "poses3D": all_poses_3d,
+            "ids": all_ids,
+            "joint_names": joint_names_3d,
+        }
+        os.makedirs(os.path.dirname(pred_export_path), exist_ok=True)
+        with open(pred_export_path, "w") as file:
+            json.dump(data, file, indent=0)
+
     # Run evaluation
     _ = evals.mpjpe.run_eval(
         labels,