[Fix] Num of tracklet does not align (#136)

1. Fix the bug mentioned in #111 caused by the number of tracklet not aligning with the total number of persons during keypoints optimization
2. Update visualization in `mview_mperson_end2end_estimator`

tools/mview_mperson_end2end_estimator.py

@@ -15,7 +15,9 @@
 from xrprimer.utils.log_utils import setup_logger
 
 from xrmocap.core.estimation.builder import build_estimator
-from xrmocap.core.visualization import visualize_project_keypoints3d
+from xrmocap.visualization.visualize_keypoints3d import (
+    visualize_keypoints3d_projected,
+)
 
 # yapf: enable
 
@@ -67,6 +69,9 @@ def main(args):
 
     # Visualization
     if not args.disable_visualization:
+        # Prepare saving path
+        if not os.path.exists(os.path.join(args.output_dir, 'kps3d')):
+            os.mkdir(os.path.join(args.output_dir, 'kps3d'))
         n_frame = args.end_frame - args.start_frame
         n_person = len(smpl_data_list)
         colors = get_different_colors(n_person)
@@ -106,12 +111,12 @@ def main(args):
                 image_list.append(image_np)
             image_array = np.array(image_list)
 
-            visualize_project_keypoints3d(
+            visualize_keypoints3d_projected(
                 keypoints=pred_keypoints3d,
-                cam_param=fisheye_param,
+                camera=fisheye_param,
                 output_path=os.path.join(args.output_dir, 'kps3d',
                                          f'project_view_{view_name}.mp4'),
-                img_arr=image_array.copy(),
+                background_arr=image_array.copy(),
                 overwrite=True)
 
             visualize_smpl_calibration(

xrmocap/transform/keypoints3d/optim/rm_duplicate.py

@@ -74,6 +74,9 @@ def optimize_keypoints3d(self, keypoints3d: Keypoints,
         n_frame, n_max_person, n_kps, _ = kps3d.shape
 
         kps3d_optim = np.full((n_frame, n_max_person, n_kps, 4), np.nan)
+        full_pid = []
+        full_keep_idxs = []
+        max_pid = 0
 
         for frame_idx in range(n_frame):
             kps3d_frame = kps3d[frame_idx, ...]
@@ -113,16 +116,33 @@ def optimize_keypoints3d(self, keypoints3d: Keypoints,
             if self.identity_tracking is not None:
                 curr_kps3d = kps3d_frame[keep_idxs, ...]
                 frame_identities = self.identity_tracking.query(curr_kps3d)
+                full_keep_idxs.append(keep_idxs)
+                full_pid.append(frame_identities)
+                if max(frame_identities) > max_pid:
+                    max_pid = max(frame_identities)
 
-                # save to Kps3d
-                kps3d_optim[frame_idx, frame_identities,
-                            ...] = kps3d_frame[keep_idxs, ...]
             else:
                 # save to Kps3d
                 n_optim_person = len(keep_idxs)
                 kps3d_optim[frame_idx, :n_optim_person,
                             ...] = kps3d_frame[keep_idxs, ...]
 
+        if self.identity_tracking is not None:
+            kps3d_optim = np.full((n_frame, max_pid + 1, n_kps, 4), np.nan)
+
+            # save to Kps3d
+            for frame_idx in range(n_frame):
+                kps3d_frame = kps3d[frame_idx, ...]
+                kps3d_frame = kps3d_frame[~np.isnan(kps3d_frame[:, 0, 0])]
+                # skip empty frame
+                if kps3d_frame.shape[0] == 0:
+                    continue
+
+                frame_identities = full_pid[frame_idx]
+                keep_idxs = full_keep_idxs[frame_idx]
+                kps3d_optim[frame_idx, frame_identities,
+                            ...] = kps3d_frame[keep_idxs, ...]
+
         keypoints3d_optim.set_keypoints(kps3d_optim)
         keypoints3d_optim.set_mask(kps3d_optim[..., -1] > 0)