vis_st4rtrack.py

"""Record3D visualizer

Parse and stream record3d captures. To get the demo data, see `./assets/download_record3d_dance.sh`.
"""

import time
from pathlib import Path

import numpy as onp
import tyro
import cv2
from tqdm.auto import tqdm

import viser
import viser.extras
import viser.transforms as tf

from glob import glob
import numpy as np
import imageio.v3 as iio
import matplotlib.pyplot as plt
import psutil

def log_memory_usage(message=""):
    """Log current memory usage with an optional message."""
    process = psutil.Process()
    memory_info = process.memory_info()
    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', conf_thre_percentile=10):
    """Load trajectory data from files.
    
    Args:
        traj_path: Path to the directory containing trajectory data
        use_float16: Whether to convert data to float16 to save memory
        max_frames: Maximum number of frames to load (None for all)
        mask_folder: Path to the directory containing mask images
        
    Returns:
        A dictionary containing loaded data
    """
    log_memory_usage("before loading data")
    
    data_cache = {
        'traj_3d_head1': None,
        'traj_3d_head2': None,
        'conf_mask_head1': None,
        'conf_mask_head2': None,
        'masks': None,
        'raw_video': None,
        'loaded': False
    }
    
    # Load masks
    masks_paths = sorted(glob(mask_folder + '/*.jpg'))
    masks = None
    
    if masks_paths:
        masks = [iio.imread(p) for p in masks_paths]
        masks = np.stack(masks, axis=0)
        # Convert masks to binary (0 or 1)
        masks = (masks < 1).astype(np.float32)
        masks = masks.sum(axis=-1) > 2  # Combine all channels, True where any channel was 1
        print(f"Original masks shape: {masks.shape}")
    else:
        print("No masks found. Will create default masks when needed.")
    
    data_cache['masks'] = masks
    
    if Path(traj_path).is_dir():
        # Find all trajectory files
        traj_3d_paths_head1 = sorted(glob(traj_path + '/pts3d1_p*.npy'), 
                                   key=lambda x: int(x.split('_p')[-1].split('.')[0]))
        conf_paths_head1 = sorted(glob(traj_path + '/conf1_p*.npy'), 
                                key=lambda x: int(x.split('_p')[-1].split('.')[0]))
        
        traj_3d_paths_head2 = sorted(glob(traj_path + '/pts3d2_p*.npy'), 
                                   key=lambda x: int(x.split('_p')[-1].split('.')[0]))
        conf_paths_head2 = sorted(glob(traj_path + '/conf2_p*.npy'), 
                                key=lambda x: int(x.split('_p')[-1].split('.')[0]))
        
        # Limit number of frames if specified
        if max_frames is not None:
            traj_3d_paths_head1 = traj_3d_paths_head1[:max_frames]
            conf_paths_head1 = conf_paths_head1[:max_frames] if conf_paths_head1 else []
            traj_3d_paths_head2 = traj_3d_paths_head2[:max_frames]
            conf_paths_head2 = conf_paths_head2[:max_frames] if conf_paths_head2 else []

        # Process head1
        if traj_3d_paths_head1:
            if use_float16:
                traj_3d_head1 = onp.stack([onp.load(p).astype(onp.float16) for p in traj_3d_paths_head1], axis=0)
            else:
                traj_3d_head1 = onp.stack([onp.load(p) for p in traj_3d_paths_head1], axis=0)
                
            log_memory_usage("after loading head1 data")
            
            h, w, _ = traj_3d_head1.shape[1:]
            num_frames = traj_3d_head1.shape[0]
            
            # If masks is None, create default masks (all ones)
            if masks is None:
                masks = np.ones((num_frames, h, w), dtype=bool)
                print(f"Created default masks with shape: {masks.shape}")
                data_cache['masks'] = masks
            else:
                # Resize masks to match trajectory dimensions using nearest neighbor interpolation
                masks_resized = np.zeros((masks.shape[0], h, w), dtype=bool)
                for i in range(masks.shape[0]):
                    masks_resized[i] = cv2.resize(
                        masks[i].astype(np.uint8), 
                        (w, h), 
                        interpolation=cv2.INTER_NEAREST
                    ).astype(bool)
                
                print(f"Resized masks shape: {masks_resized.shape}")
                data_cache['masks'] = masks_resized
            
            # Reshape trajectory data
            traj_3d_head1 = traj_3d_head1.reshape(traj_3d_head1.shape[0], -1, 6)
            data_cache['traj_3d_head1'] = traj_3d_head1
            
            if conf_paths_head1:
                conf_head1 = onp.stack([onp.load(p).astype(onp.float16) for p in conf_paths_head1], axis=0)
                conf_head1 = conf_head1.reshape(conf_head1.shape[0], -1)
                conf_head1 = conf_head1.mean(axis=0)
                # 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), conf_thre_percentile)  # Default percentile
                conf_mask_head1 = conf_head1 > conf_thre
                data_cache['conf_mask_head1'] = conf_mask_head1

        # Process head2
        if traj_3d_paths_head2:
            if use_float16:
                traj_3d_head2 = onp.stack([onp.load(p).astype(onp.float16) for p in traj_3d_paths_head2], axis=0)
            else:
                traj_3d_head2 = onp.stack([onp.load(p) for p in traj_3d_paths_head2], axis=0)
                
            log_memory_usage("after loading head2 data")
            
            # Store raw video data
            raw_video = traj_3d_head2[:, :, :, 3:6]  # [num_frames, h, w, 3]
            data_cache['raw_video'] = raw_video
            
            traj_3d_head2 = traj_3d_head2.reshape(traj_3d_head2.shape[0], -1, 6)
            data_cache['traj_3d_head2'] = traj_3d_head2
            
            if conf_paths_head2:
                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), conf_thre_percentile, axis=1)
                conf_mask_head2 = conf_head2 > conf_thre[:, None]
                data_cache['conf_mask_head2'] = conf_mask_head2
    
    data_cache['loaded'] = True
    log_memory_usage("after loading all data")
    return data_cache

def visualize_st4rtrack(
    traj_path: str = "results",
    up_dir: str = "-z", # should be +z or -z
    max_frames: int = 100,
    share: bool = False,
    point_size: float = 0.005,
    downsample_factor: int = 3,
    num_traj_points: int = 100,
    conf_thre_percentile: float = 1,
    traj_end_frame: int = 100,
    traj_start_frame: int = 0,
    traj_line_width: float = 3.,
    fixed_length_traj: int = 20,
    server: viser.ViserServer = None,
    use_float16: bool = True,
    preloaded_data: dict = None,  # Add this parameter to accept preloaded data
    color_code: str = "jet",
    # Updated hex colors: #002676 for blue and #FDB515 for red/gold
    blue_rgb: tuple[float, float, float] = (0.0, 0.149, 0.463),  # #002676
    red_rgb: tuple[float, float, float] = (0.769, 0.510, 0.055),   # #FDB515
    blend_ratio: float = 0.7,
    mask_folder: str = None,
    mid_anchor: bool = False,
    video_width: int = 320,   # Video display width
    video_height: int = 180,  # Video display height
    camera_position: tuple[float, float, float] = (1e-3, 1.5, -0.2),
) -> None:
    log_memory_usage("at start of visualization")
    
    if server is None:
        server = viser.ViserServer()
    if share:
        server.request_share_url()

    @server.on_client_connect
    def _(client: viser.ClientHandle) -> None:
        client.camera.position = camera_position
        client.camera.look_at = (0, 0, 0)

    # Configure the GUI panel size and layout
    server.gui.configure_theme(
        control_layout="collapsible",
        control_width="small",
        dark_mode=False,
        show_logo=False,
        show_share_button=True
    )

    # Add video preview to the GUI panel - placed at the top
    video_preview = server.gui.add_image(
        np.zeros((video_height, video_width, 3), dtype=np.uint8),  # Initial blank image
        format="jpeg"
    )
    
    # Use preloaded data if available
    if preloaded_data and preloaded_data.get('loaded', False):
        traj_3d_head1 = preloaded_data.get('traj_3d_head1')
        traj_3d_head2 = preloaded_data.get('traj_3d_head2')
        conf_mask_head1 = preloaded_data.get('conf_mask_head1')
        conf_mask_head2 = preloaded_data.get('conf_mask_head2')
        masks = preloaded_data.get('masks')
        raw_video = preloaded_data.get('raw_video')
        print("Using preloaded data!")
    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, 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')
        conf_mask_head2 = data.get('conf_mask_head2')
        masks = data.get('masks')
        raw_video = data.get('raw_video')

    def process_video_frame(frame_idx):
        if raw_video is None:
            return np.zeros((video_height, video_width, 3), dtype=np.uint8)
        
        # Get the original frame
        raw_frame = raw_video[frame_idx]
        
        # Adjust value range to 0-255
        if raw_frame.max() <= 1.0:
            frame = (raw_frame * 255).astype(np.uint8)
        else:
            frame = raw_frame.astype(np.uint8)
        
        # Resize to fit the preview window
        h, w = frame.shape[:2]
        # Calculate size while maintaining aspect ratio
        if h/w > video_height/video_width:  # Height limited
            new_h = video_height
            new_w = int(w * (new_h / h))
        else:  # Width limited
            new_w = video_width
            new_h = int(h * (new_w / w))
        
        # Resize
        resized_frame = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA)
        
        # Create a black background
        display_frame = np.zeros((video_height, video_width, 3), dtype=np.uint8)
        
        # Place the resized frame in the center
        y_offset = (video_height - new_h) // 2
        x_offset = (video_width - new_w) // 2
        display_frame[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = resized_frame
        
        return display_frame

    server.scene.set_up_direction(up_dir)
    print("Setting up visualization!")

    # Add visualization controls
    with server.gui.add_folder("Visualization"):
        gui_show_head1 = server.gui.add_checkbox("Tracking Points", True)
        gui_show_head2 = server.gui.add_checkbox("Recon Points", True)
        gui_show_trajectories = server.gui.add_checkbox("Trajectories", True)
        gui_use_color_tint = server.gui.add_checkbox("Use Color Tint", True)

    # Process and center point clouds
    center_point = None
    if traj_3d_head1 is not None:
        xyz_head1 = traj_3d_head1[:, :, :3]
        rgb_head1 = traj_3d_head1[:, :, 3:6]
        if center_point is None:
            center_point = onp.mean(xyz_head1, axis=(0, 1), keepdims=True)
        xyz_head1 -= center_point
        if rgb_head1.sum(axis=(-1)).max() > 125:
            rgb_head1 /= 255.0

    if traj_3d_head2 is not None:
        xyz_head2 = traj_3d_head2[:, :, :3]
        rgb_head2 = traj_3d_head2[:, :, 3:6]
        if center_point is None:
            center_point = onp.mean(xyz_head2, axis=(0, 1), keepdims=True)
        xyz_head2 -= center_point
        if rgb_head2.sum(axis=(-1)).max() > 125:
            rgb_head2 /= 255.0

    # Determine number of frames
    F = max(
        traj_3d_head1.shape[0] if traj_3d_head1 is not None else 0,
        traj_3d_head2.shape[0] if traj_3d_head2 is not None else 0
    )
    num_frames = min(max_frames, F)
    traj_end_frame = min(traj_end_frame, num_frames)
    print(f"Number of frames: {num_frames}")
    xyz_head1 = xyz_head1[:num_frames]
    xyz_head2 = xyz_head2[:num_frames]
    rgb_head1 = rgb_head1[:num_frames]
    rgb_head2 = rgb_head2[:num_frames]

    # Add playback UI.
    with server.gui.add_folder("Playback"):
        gui_timestep = server.gui.add_slider(
            "Timestep",
            min=0,
            max=num_frames - 1,
            step=1,
            initial_value=0,
            disabled=True,
        )
        gui_next_frame = server.gui.add_button("Next Frame", disabled=True)
        gui_prev_frame = server.gui.add_button("Prev Frame", disabled=True)
        gui_playing = server.gui.add_checkbox("Playing", True)
        gui_framerate = server.gui.add_slider(
            "FPS", min=1, max=60, step=0.1, initial_value=20
        )
        gui_framerate_options = server.gui.add_button_group(
            "FPS options", ("10", "20", "30")
        )
        gui_show_all_frames = server.gui.add_checkbox("Show all frames", False)
        gui_stride = server.gui.add_slider(
            "Stride",
            min=1,
            max=num_frames,
            step=1,
            initial_value=5,
            disabled=True,  # Initially disabled
        )

    # Frame step buttons.
    @gui_next_frame.on_click
    def _(_) -> None:
        gui_timestep.value = (gui_timestep.value + 1) % num_frames

    @gui_prev_frame.on_click
    def _(_) -> None:
        gui_timestep.value = (gui_timestep.value - 1) % num_frames

    # Disable frame controls when we're playing.
    @gui_playing.on_update
    def _(_) -> None:
        gui_timestep.disabled = gui_playing.value or gui_show_all_frames.value
        gui_next_frame.disabled = gui_playing.value or gui_show_all_frames.value
        gui_prev_frame.disabled = gui_playing.value or gui_show_all_frames.value

    # Set the framerate when we click one of the options.
    @gui_framerate_options.on_click
    def _(_) -> None:
        gui_framerate.value = int(gui_framerate_options.value)

    prev_timestep = gui_timestep.value

    # Toggle frame visibility when the timestep slider changes.
    @gui_timestep.on_update
    def _(_) -> None:
        nonlocal prev_timestep
        current_timestep = gui_timestep.value
        if not gui_show_all_frames.value:
            with server.atomic():
                if gui_show_head1.value:
                    frame_nodes_head1[current_timestep].visible = True
                    frame_nodes_head1[prev_timestep].visible = False
                if gui_show_head2.value:
                    frame_nodes_head2[current_timestep].visible = True
                    frame_nodes_head2[prev_timestep].visible = False
        prev_timestep = current_timestep
        server.flush()  # Optional!

    # Show or hide all frames based on the checkbox.
    @gui_show_all_frames.on_update
    def _(_) -> None:
        gui_stride.disabled = not gui_show_all_frames.value  # Enable/disable stride slider
        if gui_show_all_frames.value:
            # Show frames with stride
            stride = gui_stride.value
            with server.atomic():
                for i, (node1, node2) in enumerate(zip(frame_nodes_head1, frame_nodes_head2)):
                    node1.visible = gui_show_head1.value and (i % stride == 0)
                    node2.visible = gui_show_head2.value and (i % stride == 0)
            # Disable playback controls
            gui_playing.disabled = True
            gui_timestep.disabled = True
            gui_next_frame.disabled = True
            gui_prev_frame.disabled = True
        else:
            # Show only the current frame
            current_timestep = gui_timestep.value
            with server.atomic():
                for i, (node1, node2) in enumerate(zip(frame_nodes_head1, frame_nodes_head2)):
                    node1.visible = gui_show_head1.value and (i == current_timestep)
                    node2.visible = gui_show_head2.value and (i == current_timestep)
            # Re-enable playback controls
            gui_playing.disabled = False
            gui_timestep.disabled = gui_playing.value
            gui_next_frame.disabled = gui_playing.value
            gui_prev_frame.disabled = gui_playing.value

    # Update frame visibility when the stride changes.
    @gui_stride.on_update
    def _(_) -> None:
        if gui_show_all_frames.value:
            # Update frame visibility based on new stride
            stride = gui_stride.value
            with server.atomic():
                for i, (node1, node2) in enumerate(zip(frame_nodes_head1, frame_nodes_head2)):
                    node1.visible = gui_show_head1.value and (i % stride == 0)
                    node2.visible = gui_show_head2.value and (i % stride == 0)

    # Load in frames.
    server.scene.add_frame(
        "/frames",
        wxyz=tf.SO3.exp(onp.array([onp.pi / 2.0, 0.0, 0.0])).wxyz,
        position=(0, 0, 0),
        show_axes=False,
    )
    frame_nodes_head1: list[viser.FrameHandle] = []
    frame_nodes_head2: list[viser.FrameHandle] = []

    # Extract RGB components for tinting
    blue_r, blue_g, blue_b = blue_rgb
    red_r, red_g, red_b = red_rgb
    
    # Create frames for each timestep
    frame_nodes_head1 = []
    frame_nodes_head2 = []
    for i in tqdm(range(num_frames)):
        # Process head1
        if traj_3d_head1 is not None:
            frame_nodes_head1.append(server.scene.add_frame(f"/frames/t{i}/head1", show_axes=False))
            position = xyz_head1[i]
            color = rgb_head1[i]
            if conf_mask_head1 is not None:
                position = position[conf_mask_head1[i]]
                color = color[conf_mask_head1[i]]
            
            # Add point cloud for head1 with optional blue tint
            color_head1 = color.copy()
            if gui_use_color_tint.value:
                color_head1 *= blend_ratio
                color_head1[:, 0] = onp.clip(color_head1[:, 0] + blue_r * (1 - blend_ratio), 0, 1)  # R
                color_head1[:, 1] = onp.clip(color_head1[:, 1] + blue_g * (1 - blend_ratio), 0, 1)  # G
                color_head1[:, 2] = onp.clip(color_head1[:, 2] + blue_b * (1 - blend_ratio), 0, 1)  # B
            
            server.scene.add_point_cloud(
                name=f"/frames/t{i}/head1/point_cloud",
                points=position[::downsample_factor],
                colors=color_head1[::downsample_factor],
                point_size=point_size,
                point_shape="rounded",
            )

        # Process head2
        if traj_3d_head2 is not None:
            frame_nodes_head2.append(server.scene.add_frame(f"/frames/t{i}/head2", show_axes=False))
            position = xyz_head2[i]
            color = rgb_head2[i]
            if conf_mask_head2 is not None:
                position = position[conf_mask_head2[i]]
                color = color[conf_mask_head2[i]]
            
            # Add point cloud for head2 with optional red tint
            color_head2 = color.copy()
            if gui_use_color_tint.value:
                color_head2 *= blend_ratio
                color_head2[:, 0] = onp.clip(color_head2[:, 0] + red_r * (1 - blend_ratio), 0, 1)  # R
                color_head2[:, 1] = onp.clip(color_head2[:, 1] + red_g * (1 - blend_ratio), 0, 1)  # G
                color_head2[:, 2] = onp.clip(color_head2[:, 2] + red_b * (1 - blend_ratio), 0, 1)  # B
            
            server.scene.add_point_cloud(
                name=f"/frames/t{i}/head2/point_cloud",
                points=position[::downsample_factor],
                colors=color_head2[::downsample_factor],
                point_size=point_size,
                point_shape="rounded",
            )

    # Update visibility based on checkboxes
    @gui_show_head1.on_update
    def _(_) -> None:
        with server.atomic():
            for frame_node in frame_nodes_head1:
                frame_node.visible = gui_show_head1.value and (
                    gui_show_all_frames.value
                    or (not gui_show_all_frames.value )
                )

    @gui_show_head2.on_update
    def _(_) -> None:
        with server.atomic():
            for frame_node in frame_nodes_head2:
                frame_node.visible = gui_show_head2.value and (
                    gui_show_all_frames.value
                    or (not gui_show_all_frames.value )
                )

    # Initial visibility
    for i, (node1, node2) in enumerate(zip(frame_nodes_head1, frame_nodes_head2)):
        if gui_show_all_frames.value:
            node1.visible = gui_show_head1.value and (i % gui_stride.value == 0)
            node2.visible = gui_show_head2.value and (i % gui_stride.value == 0)
        else:
            node1.visible = gui_show_head1.value and (i == gui_timestep.value)
            node2.visible = gui_show_head2.value and (i == gui_timestep.value)

    # Process and visualize trajectories for head1
    if traj_3d_head1 is not None:
        # Get points over time
        xyz_head1_centered = xyz_head1.copy()
        
        # Select points to visualize
        num_points = xyz_head1.shape[1]
        points_to_visualize = min(num_points, num_traj_points)
        
        # Get the mask for the first frame and reshape it to match point cloud dimensions
        if mid_anchor:
            first_frame_mask = masks[num_frames//2].reshape(-1)
        else:
            first_frame_mask = masks[0].reshape(-1) #[#points, h]
        
        # Calculate trajectory lengths for each point
        trajectories = xyz_head1_centered[traj_start_frame:traj_end_frame]  # Shape: (num_frames, num_points, 3)
        traj_diffs = np.diff(trajectories, axis=0)  # Differences between consecutive frames
        traj_lengths = np.sum(np.sqrt(np.sum(traj_diffs**2, axis=-1)), axis=0)  # Sum of distances for each point
        
        # Get points that are within the mask
        valid_indices = np.where(first_frame_mask)[0]
        
        if len(valid_indices) > 0:
            # Calculate average trajectory length for masked points
            masked_traj_lengths = traj_lengths[valid_indices]
            avg_traj_length = np.mean(masked_traj_lengths)
            
            if mask_folder is not None:
                # do not filter points by trajectory length
                long_traj_indices = valid_indices
            else:
                # Filter points by trajectory length
                long_traj_indices = valid_indices[masked_traj_lengths >= avg_traj_length]
            
            # Randomly sample from the filtered points
            if len(long_traj_indices) > 0:
                # Random sampling without replacement
                selected_indices = np.random.choice(
                    len(long_traj_indices),
                    min(points_to_visualize, len(long_traj_indices)),
                    replace=False
                )
                # Get the actual indices in their original order
                valid_point_indices = long_traj_indices[np.sort(selected_indices)]
            else:
                valid_point_indices = np.array([])
        else:
            valid_point_indices = np.array([])
        
        if len(valid_point_indices) > 0:
            # Get trajectories for all valid points
            trajectories = xyz_head1_centered[traj_start_frame:traj_end_frame, valid_point_indices]
            N_point = trajectories.shape[1]
            if color_code == "rainbow":
                point_colors = plt.cm.rainbow(np.linspace(0, 1, N_point))[:, :3]
            elif color_code == "jet":
                point_colors = plt.cm.jet(np.linspace(0, 1, N_point))[:, :3]
            # Modify the loop to handle frames less than fixed_length_traj
            for i in range(traj_end_frame - traj_start_frame):
                # Calculate the actual trajectory length for this frame
                actual_length = min(fixed_length_traj, i + 1)
                
                if actual_length > 1:  # Need at least 2 points to form a line
                    # Get the appropriate slice of trajectory data
                    start_idx = max(0, i - actual_length + 1)
                    end_idx = i + 1
                    
                    # Create line segments between consecutive frames
                    traj_slice = trajectories[start_idx:end_idx]
                    line_points = np.stack([traj_slice[:-1], traj_slice[1:]], axis=2)
                    line_points = line_points.reshape(-1, 2, 3)
                    
                    # Create corresponding colors
                    line_colors = np.tile(point_colors, (actual_length-1, 1))
                    line_colors = np.stack([line_colors, line_colors], axis=1)
                    
                    # Add line segments
                    server.scene.add_line_segments(
                        name=f"/frames/t{i+traj_start_frame}/head1/trajectory",
                        points=line_points,
                        colors=line_colors,
                        line_width=traj_line_width,
                        visible=gui_show_trajectories.value
                    )

    # Add trajectory controls functionality
    @gui_show_trajectories.on_update
    def _(_) -> None:
        with server.atomic():
            # Remove all existing trajectories
            for i in range(num_frames):
                try:
                    server.scene.remove_by_name(f"/frames/t{i}/head1/trajectory")
                except KeyError:
                    pass
            
            # Create new trajectories if enabled
            if gui_show_trajectories.value and traj_3d_head1 is not None:
                # Get the mask for the last frame and reshape it
                last_frame_mask = masks[traj_end_frame-1].reshape(-1)
                
                # Calculate trajectory lengths
                trajectories = xyz_head1_centered[traj_start_frame:traj_end_frame]
                traj_diffs = np.diff(trajectories, axis=0)
                traj_lengths = np.sum(np.sqrt(np.sum(traj_diffs**2, axis=-1)), axis=0)
                
                # Get points that are within the mask
                valid_indices = np.where(last_frame_mask)[0]
                
                if len(valid_indices) > 0:
                    # Filter by trajectory length
                    masked_traj_lengths = traj_lengths[valid_indices]
                    avg_traj_length = np.mean(masked_traj_lengths)
                    long_traj_indices = valid_indices[masked_traj_lengths >= avg_traj_length]
                    
                    # Randomly sample from the filtered points
                    if len(long_traj_indices) > 0:
                        # Random sampling without replacement
                        selected_indices = np.random.choice(
                            len(long_traj_indices),
                            min(points_to_visualize, len(long_traj_indices)),
                            replace=False
                        )
                        # Get the actual indices in their original order
                        valid_point_indices = long_traj_indices[np.sort(selected_indices)]
                    else:
                        valid_point_indices = np.array([])
                else:
                    valid_point_indices = np.array([])

                if len(valid_point_indices) > 0:
                    # Get trajectories for all valid points
                    trajectories = xyz_head1_centered[traj_start_frame:traj_end_frame, valid_point_indices]
                    N_point = trajectories.shape[1]

                    if color_code == "rainbow":
                        point_colors = plt.cm.rainbow(np.linspace(0, 1, N_point))[:, :3]
                    elif color_code == "jet":
                        point_colors = plt.cm.jet(np.linspace(0, 1, N_point))[:, :3]
                    
                    # Modify the loop to handle frames less than fixed_length_traj
                    for i in range(traj_end_frame - traj_start_frame):
                        # Calculate the actual trajectory length for this frame
                        actual_length = min(fixed_length_traj, i + 1)
                        
                        if actual_length > 1:  # Need at least 2 points to form a line
                            # Get the appropriate slice of trajectory data
                            start_idx = max(0, i - actual_length + 1)
                            end_idx = i + 1
                            
                            # Create line segments between consecutive frames
                            traj_slice = trajectories[start_idx:end_idx]
                            line_points = np.stack([traj_slice[:-1], traj_slice[1:]], axis=2)
                            line_points = line_points.reshape(-1, 2, 3)
                            
                            # Create corresponding colors
                            line_colors = np.tile(point_colors, (actual_length-1, 1))
                            line_colors = np.stack([line_colors, line_colors], axis=1)
                            
                            # Add line segments
                            server.scene.add_line_segments(
                                name=f"/frames/t{i+traj_start_frame}/head1/trajectory",
                                points=line_points,
                                colors=line_colors,
                                line_width=traj_line_width,
                                visible=True
                            )

    # Update color tinting when the checkbox changes
    @gui_use_color_tint.on_update
    def _(_) -> None:
        with server.atomic():
            for i in range(num_frames):
                # Update head1 point cloud
                if traj_3d_head1 is not None:
                    position = xyz_head1[i]
                    color = rgb_head1[i]
                    if conf_mask_head1 is not None:
                        position = position[conf_mask_head1[i]]
                        color = color[conf_mask_head1[i]]
                    
                    color_head1 = color.copy()
                    if gui_use_color_tint.value:
                        color_head1 *= blend_ratio
                        color_head1[:, 0] = onp.clip(color_head1[:, 0] + blue_r * (1 - blend_ratio), 0, 1)  # R
                        color_head1[:, 1] = onp.clip(color_head1[:, 1] + blue_g * (1 - blend_ratio), 0, 1)  # G
                        color_head1[:, 2] = onp.clip(color_head1[:, 2] + blue_b * (1 - blend_ratio), 0, 1)  # B
                    
                    server.scene.remove_by_name(f"/frames/t{i}/head1/point_cloud")
                    server.scene.add_point_cloud(
                        name=f"/frames/t{i}/head1/point_cloud",
                        points=position[::downsample_factor],
                        colors=color_head1[::downsample_factor],
                        point_size=point_size,
                        point_shape="rounded",
                    )
                
                # Update head2 point cloud
                if traj_3d_head2 is not None:
                    position = xyz_head2[i]
                    color = rgb_head2[i]
                    if conf_mask_head2 is not None:
                        position = position[conf_mask_head2[i]]
                        color = color[conf_mask_head2[i]]
                    
                    color_head2 = color.copy()
                    if gui_use_color_tint.value:
                        color_head2 *= blend_ratio
                        color_head2[:, 0] = onp.clip(color_head2[:, 0] + red_r * (1 - blend_ratio), 0, 1)  # R
                        color_head2[:, 1] = onp.clip(color_head2[:, 1] + red_g * (1 - blend_ratio), 0, 1)  # G
                        color_head2[:, 2] = onp.clip(color_head2[:, 2] + red_b * (1 - blend_ratio), 0, 1)  # B
                    
                    server.scene.remove_by_name(f"/frames/t{i}/head2/point_cloud")
                    server.scene.add_point_cloud(
                        name=f"/frames/t{i}/head2/point_cloud",
                        points=position[::downsample_factor],
                        colors=color_head2[::downsample_factor],
                        point_size=point_size,
                        point_shape="rounded",
                    )

    # Initialize video preview
    if raw_video is not None:
        video_preview.image = process_video_frame(0)
    
    # Update video preview when timestep changes
    @gui_timestep.on_update
    def _(_) -> None:
        current_timestep = gui_timestep.value
        if raw_video is not None:
            video_preview.image = process_video_frame(current_timestep)
    
    # Playback update loop.
    log_memory_usage("before starting playback loop")
    
    prev_timestep = gui_timestep.value
    while True:
        current_timestep = gui_timestep.value
        
        # If timestep changes, update frame visibility
        if current_timestep != prev_timestep:
            with server.atomic():
                # ... existing code ...
                
                # Update video preview
                if raw_video is not None:
                    video_preview.image = process_video_frame(current_timestep)
        
        # Update in playback mode
        if gui_playing.value and not gui_show_all_frames.value:
            gui_timestep.value = (gui_timestep.value + 1) % num_frames
            
            # Update video preview in playback mode
            if raw_video is not None:
                video_preview.image = process_video_frame(gui_timestep.value)
        
        time.sleep(1.0 / gui_framerate.value)


if __name__ == "__main__":
    tyro.cli(visualize_st4rtrack)