From 53f13d6f97f4f48df3ef22d91da05d91ce5cc1e2 Mon Sep 17 00:00:00 2001
From: Daniel <daniel@mail.de>
Date: Thu, 12 Jun 2025 12:43:09 +0200
Subject: [PATCH] Some mixed updates.

---
 .../rpt3d_wrapper_cpp/src/rpt3d_wrapper.cpp   | 14 ++--
 scripts/test_skelda_dataset.py                | 64 +++++++++++++++++--
 2 files changed, 67 insertions(+), 11 deletions(-)

diff --git a/extras/ros/rpt3d_wrapper_cpp/src/rpt3d_wrapper.cpp b/extras/ros/rpt3d_wrapper_cpp/src/rpt3d_wrapper.cpp
index 9a17172..e99641a 100644
--- a/extras/ros/rpt3d_wrapper_cpp/src/rpt3d_wrapper.cpp
+++ b/extras/ros/rpt3d_wrapper_cpp/src/rpt3d_wrapper.cpp
@@ -251,13 +251,13 @@ void Rpt3DWrapperNode::call_model()
     {
         auto pose_tracks = pose_tracker->track_poses(poses_3d, joint_names, min_ts);
         std::vector<std::vector<std::array<float, 4>>> poses_3d_refined;
-    for (size_t j = 0; j < pose_tracks.size(); j++)
-    {
-        auto &pose = std::get<1>(pose_tracks[j]);
-        poses_3d_refined.push_back(pose);
-        auto &track_id = std::get<0>(pose_tracks[j]);
-        track_ids.push_back(track_id);
-    }
+        for (size_t j = 0; j < pose_tracks.size(); j++)
+        {
+            auto &pose = std::get<1>(pose_tracks[j]);
+            poses_3d_refined.push_back(pose);
+            auto &track_id = std::get<0>(pose_tracks[j]);
+            track_ids.push_back(track_id);
+        }
         valid_poses = std::move(poses_3d_refined);
     }
     else
diff --git a/scripts/test_skelda_dataset.py b/scripts/test_skelda_dataset.py
index 688bb52..933cc4b 100644
--- a/scripts/test_skelda_dataset.py
+++ b/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,