refactor the code

app.py

@@ -1,132 +1,15 @@
 import random
 import threading
 import psutil
-import numpy as onp
 import fastapi
 import gradio as gr
 import uvicorn
-import os
-from pathlib import Path
-from glob import glob
-import cv2
-import numpy as np
-import imageio.v3 as iio
 
 from viser_proxy_manager import ViserProxyManager
-from vis_st4rtrack import visualize_st4rtrack, log_memory_usage
+from vis_st4rtrack import visualize_st4rtrack, load_trajectory_data, log_memory_usage
 
 # Global cache for loaded data
-global_data_cache = {
-    'traj_3d_head1': None,
-    'traj_3d_head2': None,
-    'conf_mask_head1': None,
-    'conf_mask_head2': None,
-    'masks': None,
-    'loaded': False
-}
-
-def load_data_once(traj_path="results", use_float16=True):
-    """Load data once and store in global cache."""
-    if global_data_cache['loaded']:
-        return
-    
-    log_memory_usage("before loading data")
-    
-    # Load masks
-    mask_folder = './train'
-    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.")
-    
-    global_data_cache['masks'] = masks
-    
-    if Path(traj_path).is_dir():
-        # Load head1 data
-        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]))
-        
-        # Load head2 data
-        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]))
-
-        # 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:]
-            
-            # If masks is None, create default masks (all ones)
-            if masks is None:
-                num_frames = traj_3d_head1.shape[0]
-                masks = np.ones((num_frames, h, w), dtype=bool)
-                print(f"Created default masks with shape: {masks.shape}")
-                global_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}")
-                global_data_cache['masks'] = masks_resized
-            
-            # Reshape trajectory data
-            traj_3d_head1 = traj_3d_head1.reshape(traj_3d_head1.shape[0], -1, 6)
-            global_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_mask_head1 = conf_head1 > 1.0  # Default threshold
-                global_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")
-            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
-            
-            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)
-                conf_mask_head2 = conf_head2 > 1.0  # Default threshold
-                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")
-
+global_data_cache = None
 
 def check_ram_usage(threshold_percent=90):
     """Check if RAM usage is above the threshold.
@@ -143,20 +26,16 @@ def check_ram_usage(threshold_percent=90):
 
 
 def main() -> None:
-    # Load data once at startup
-    load_data_once(use_float16=True)
+    # Load data once at startup using the function from vis_st4rtrack.py
+    global global_data_cache
+    global_data_cache = load_trajectory_data(use_float16=True, max_frames=32)
     
     app = fastapi.FastAPI()
     viser_manager = ViserProxyManager(app)
 
-
     # Create a Gradio interface with title, iframe, and buttons
     with gr.Blocks(title="Viser Viewer") as demo:
-        # Add a title and description
-        # gr.Markdown("# 🌐 Viser Interactive Viewer Test")
-
         # Add the iframe with a border
-        # add_sphere_btn = gr.Button("Add Random Sphere")
         iframe_html = gr.HTML("")
         status_text = gr.Markdown("")  # Add status text component
 
@@ -194,13 +73,8 @@ def main() -> None:
                     "use_float16": True,
                     "preloaded_data": global_data_cache,  # Pass the preloaded data
                     "color_code": "jet",
-                    # "blue_rgb": (0.22, 0.82, 1.0),  # #37D2FF
-                    # "red_rgb": (1.0, 0.39, 0.22),   # #FF6337
                     "blue_rgb": (0.0, 0.149, 0.463),  # #002676
                     "red_rgb": (0.769, 0.510, 0.055),   # #FDB515
-                    # "color_code": "rainbow",
-                    # "blue_rgb": (0., 0., 1.),
-                    # "red_rgb": (1., 0., 0.),
                     "blend_ratio": 0.7
                 },
                 daemon=True
@@ -217,7 +91,6 @@ def main() -> None:
                 loading="lazy"
             ></iframe>
             """, "**System Status:** Visualization loaded successfully."
-        
 
         @demo.unload
         def stop(request: gr.Request):

vis_st4rtrack.py

@@ -28,6 +28,138 @@ 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'):
+    """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), 1)  # 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), 1, 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
@@ -81,7 +213,26 @@ def visualize_st4rtrack(
         format="jpeg"
     )
     
-    # Create a function to process video frames and resize them
+    # 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)
+        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)
@@ -121,106 +272,6 @@ def visualize_st4rtrack(
     server.scene.set_up_direction(up_dir)
     print("Setting up visualization!")
 
-    # Use preloaded data if available
-    if preloaded_data and preloaded_data['loaded']:
-        traj_3d_head1 = preloaded_data['traj_3d_head1']
-        traj_3d_head2 = preloaded_data['traj_3d_head2']
-        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)
-        print("No preloaded data available, loading from files...")
-        # Load both head1 and head2 data
-        traj_3d_head1 = None
-        traj_3d_head2 = None
-        conf_mask_head1 = None
-        conf_mask_head2 = None
-        masks = None
-        if mask_folder is not None:
-            masks_paths = sorted(glob(mask_folder + '/*.jpg'))
-            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}")
-        
-        if Path(traj_path).is_dir():
-            # Load head1 data
-            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]))
-            
-            # Load head2 data
-            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]))
-
-            # 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}")
-                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}")
-                    masks = masks_resized
-                
-                # Reshape trajectory data
-                traj_3d_head1 = traj_3d_head1.reshape(traj_3d_head1.shape[0], -1, 6)
-                
-                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)
-                    conf_mask_head1 = conf_head1 > conf_thre
-
-            # 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")
-                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)
-                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 10 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]
-
     # Add visualization controls
     with server.gui.add_folder("Visualization"):
         gui_show_head1 = server.gui.add_checkbox("Tracking Points", True)
@@ -286,7 +337,7 @@ def visualize_st4rtrack(
             min=1,
             max=num_frames,
             step=1,
-            initial_value=1,
+            initial_value=5,
             disabled=True,  # Initially disabled
         )