minor bug

app.py

@@ -112,6 +112,7 @@ def load_data_once(traj_path="results", use_float16=True):
                 traj_3d_head2 = onp.stack([onp.load(p) for p in traj_3d_paths_head2], axis=0)
                 
             log_memory_usage("after loading head2 data")
+            raw_video = traj_3d_head2[:, :, :, 3:6] # [num_frames, h, w, 3]
             
             traj_3d_head2 = traj_3d_head2.reshape(traj_3d_head2.shape[0], -1, 6)
             global_data_cache['traj_3d_head2'] = traj_3d_head2
@@ -123,6 +124,7 @@ def load_data_once(traj_path="results", use_float16=True):
                 global_data_cache['conf_mask_head2'] = conf_mask_head2
     
     global_data_cache['loaded'] = True
+    global_data_cache['raw_video'] = raw_video
     log_memory_usage("after loading all data")
 
 

vis_st4rtrack.py

@@ -128,6 +128,7 @@ def visualize_st4rtrack(
         conf_mask_head1 = preloaded_data['conf_mask_head1']
         conf_mask_head2 = preloaded_data['conf_mask_head2']
         masks = preloaded_data['masks']
+        raw_video = preloaded_data['raw_video']
         print("Using preloaded data!")
     else:
         # Original data loading code (as a fallback)