minor

vis_st4rtrack.py

@@ -28,7 +28,7 @@ def log_memory_usage(message=""):
     memory_mb = memory_info.rss / (1024 * 1024)  # Convert to MB
     print(f"Memory usage {message}: {memory_mb:.2f} MB")
 
-def load_trajectory_data(traj_path="results", use_float16=True, max_frames=None, mask_folder='./train'):
+def load_trajectory_data(traj_path="results", use_float16=True, max_frames=None, mask_folder='./train', conf_thre_percentile=10):
     """Load trajectory data from files.
     
     Args:
@@ -128,7 +128,7 @@ def load_trajectory_data(traj_path="results", use_float16=True, max_frames=None,
                 # repeat the conf_head1 to match the number of frames in the dimension 0
                 conf_head1 = np.tile(conf_head1, (num_frames, 1))
                 # Convert to float32 before calculating percentile to avoid overflow
-                conf_thre = np.percentile(conf_head1.astype(np.float32), 1)  # Default percentile
+                conf_thre = np.percentile(conf_head1.astype(np.float32), conf_thre_percentile)  # Default percentile
                 conf_mask_head1 = conf_head1 > conf_thre
                 data_cache['conf_mask_head1'] = conf_mask_head1
 
@@ -152,7 +152,7 @@ def load_trajectory_data(traj_path="results", use_float16=True, max_frames=None,
                 conf_head2 = onp.stack([onp.load(p).astype(onp.float16) for p in conf_paths_head2], axis=0)
                 conf_head2 = conf_head2.reshape(conf_head2.shape[0], -1)
                 # set conf thre to be 1 percentile of the conf_head2, for each frame
-                conf_thre = np.percentile(conf_head2.astype(np.float32), 1, axis=1)
+                conf_thre = np.percentile(conf_head2.astype(np.float32), conf_thre_percentile, axis=1)
                 conf_mask_head2 = conf_head2 > conf_thre[:, None]
                 data_cache['conf_mask_head2'] = conf_mask_head2
     
@@ -165,7 +165,7 @@ def visualize_st4rtrack(
     up_dir: str = "-z", # should be +z or -z
     max_frames: int = 100,
     share: bool = False,
-    point_size: float = 0.006,
+    point_size: float = 0.005,
     downsample_factor: int = 3,
     num_traj_points: int = 100,
     conf_thre_percentile: float = 1,
@@ -226,7 +226,7 @@ def visualize_st4rtrack(
     else:
         # Load data using the shared function
         print("No preloaded data available, loading from files...")
-        data = load_trajectory_data(traj_path, use_float16, max_frames, mask_folder)
+        data = load_trajectory_data(traj_path, use_float16, max_frames, mask_folder, conf_thre_percentile)
         traj_3d_head1 = data.get('traj_3d_head1')
         traj_3d_head2 = data.get('traj_3d_head2')
         conf_mask_head1 = data.get('conf_mask_head1')