refactor: enhance track visualization functionality

app.py

@@ -24,26 +24,22 @@ from models.SpaTrackV2.models.vggt4track.models.vggt_moe import VGGT4Track
 from models.SpaTrackV2.models.vggt4track.utils.load_fn import preprocess_image
 from models.SpaTrackV2.models.predictor import Predictor
 from models.SpaTrackV2.models.utils import get_points_on_a_grid
+from matplotlib import cm
 
-
-# --- TTM SPECIFIC IMPORTS ---
 from diffusers.utils import export_to_video, load_image
-# Note: Ensure pipelines/wan_pipeline.py and pipelines/utils.py are in your working directory
+
 from pipelines.wan_pipeline import WanImageToVideoTTMPipeline
-from pipelines.utils import compute_hw_from_area, validate_inputs
+from pipelines.utils import compute_hw_from_area
 
-# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Constants
 MAX_FRAMES = 81
 OUTPUT_FPS = 24
 RENDER_WIDTH = 512
 RENDER_HEIGHT = 384
 WAN_MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
 
-# Camera movement types
 CAMERA_MOVEMENTS = [
     "static",
     "move_forward",
@@ -54,7 +50,6 @@ CAMERA_MOVEMENTS = [
     "move_down"
 ]
 
-# Thread pool for delayed deletion
 thread_pool_executor = ThreadPoolExecutor(max_workers=2)
 
 
@@ -84,7 +79,6 @@ def create_user_temp_dir():
     return temp_dir
 
 
-# Global model initialization for Spatial Tracker
 print("🚀 Initializing tracking models...")
 
 vggt4track_model = VGGT4Track.from_pretrained(
@@ -200,6 +194,145 @@ def render_from_pointcloud(rgb_frames, depth_frames, intrinsics, original_extrin
     return {'rendered': output_path, 'motion_signal': motion_signal_path, 'mask': mask_path}
 
 
+def visualize_tracks_with_selective_history(
+    video: np.ndarray,  # (T, H, W, C) uint8
+    tracks: np.ndarray,  # (T, N, 2) float - 2D track coordinates
+    track_depths: np.ndarray,  # (T, N) float - depth at each track point
+    visibility: np.ndarray = None,  # (T, N) float
+    output_path: str = None,
+    fps: int = 24,
+    depth_threshold: float = None,  # Threshold to separate foreground/background
+    bg_history_length: int = -1,  # History length for background (-1 = infinite)
+    fg_history_length: int = 0,  # History length for foreground (0 = no history)
+    linewidth: int = 2,
+    point_radius: int = 4,
+):
+    """
+    Visualize tracked points with selective history trails.
+
+    Background points (depth > threshold): Show history trails
+    Foreground points (depth <= threshold): No history trails (only current position)
+
+    Args:
+        video: Input video frames (T, H, W, C)
+        tracks: 2D track positions (T, N, 2) - (x, y) coordinates
+        track_depths: Depth values at each track point (T, N)
+        visibility: Visibility mask for each track point (T, N)
+        output_path: Path to save the output video
+        fps: Output video frame rate
+        depth_threshold: Depth threshold to separate FG/BG. If None, uses median.
+        bg_history_length: History length for background tracks (-1 = infinite)
+        fg_history_length: History length for foreground tracks (0 = no history)
+        linewidth: Width of the history trail lines
+        point_radius: Radius of the current position marker
+
+    Returns:
+        Annotated video frames (T, H, W, C)
+    """
+    T, H, W, C = video.shape
+    N = tracks.shape[1]
+
+    # Compute depth threshold if not provided
+    # Use median depth at frame 0 as the threshold
+    if depth_threshold is None:
+        valid_depths = track_depths[0][track_depths[0] > 0]
+        if len(valid_depths) > 0:
+            depth_threshold = np.median(valid_depths)
+        else:
+            depth_threshold = 1.0
+
+    # Classify points as foreground or background based on initial depth
+    # Foreground = closer (smaller depth), Background = farther (larger depth)
+    initial_depths = track_depths[0]
+    is_background = initial_depths > depth_threshold
+
+    # Create color palette - rainbow based on y-coordinate
+    color_map = cm.get_cmap("gist_rainbow")
+    y_min, y_max = tracks[0, :, 1].min(), tracks[0, :, 1].max()
+    if y_max - y_min < 1e-6:
+        y_max = y_min + 1
+
+    colors = np.zeros((N, 3), dtype=np.uint8)
+    for i in range(N):
+        norm_y = (tracks[0, i, 1] - y_min) / (y_max - y_min)
+        color = np.array(color_map(norm_y)[:3]) * 255
+        colors[i] = color.astype(np.uint8)
+
+    # Process each frame
+    res_video = video.copy()
+
+    for t in range(T):
+        frame = res_video[t].copy()
+
+        # Draw history trails
+        # Background points: draw full history
+        # Foreground points: draw limited or no history
+        for i in range(N):
+            # Determine history length based on foreground/background
+            if is_background[i]:
+                history_len = bg_history_length
+            else:
+                history_len = fg_history_length
+
+            # Calculate start frame for history
+            if history_len < 0:  # Infinite history
+                start_frame = 0
+            elif history_len == 0:  # No history
+                start_frame = t
+            else:
+                start_frame = max(0, t - history_len)
+
+            # Draw history trail (lines connecting past positions)
+            if start_frame < t:
+                for j in range(start_frame, t):
+                    # Check visibility
+                    if visibility is not None:
+                        if visibility[j, i] < 0.5 or visibility[j + 1, i] < 0.5:
+                            continue
+
+                    pt1 = (int(tracks[j, i, 0]), int(tracks[j, i, 1]))
+                    pt2 = (int(tracks[j + 1, i, 0]), int(tracks[j + 1, i, 1]))
+
+                    # Skip invalid coordinates
+                    if pt1[0] <= 0 or pt1[1] <= 0 or pt2[0] <= 0 or pt2[1] <= 0:
+                        continue
+                    if pt1[0] >= W or pt1[1] >= H or pt2[0] >= W or pt2[1] >= H:
+                        continue
+
+                    # Draw with fading alpha (older = more transparent)
+                    alpha = (j - start_frame + 1) / (t - start_frame + 1)
+                    color = colors[i].tolist()
+                    cv2.line(frame, pt1, pt2, color, linewidth, cv2.LINE_AA)
+
+        # Draw current positions (dots)
+        for i in range(N):
+            if visibility is not None and visibility[t, i] < 0.5:
+                continue
+
+            coord = (int(tracks[t, i, 0]), int(tracks[t, i, 1]))
+
+            # Skip invalid coordinates
+            if coord[0] <= 0 or coord[1] <= 0:
+                continue
+            if coord[0] >= W or coord[1] >= H:
+                continue
+
+            color = colors[i].tolist()
+            cv2.circle(frame, coord, point_radius, color, -1)
+
+        res_video[t] = frame
+
+    # Save video if output path provided
+    if output_path:
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        out = cv2.VideoWriter(output_path, fourcc, fps, (W, H))
+        for frame in res_video:
+            out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+        out.release()
+
+    return res_video
+
+
 @spaces.GPU
 def run_spatial_tracker(video_tensor: torch.Tensor):
     """
@@ -211,7 +344,6 @@ def run_spatial_tracker(video_tensor: torch.Tensor):
     Returns:
         Dictionary containing tracking results
     """
-    # Run VGGT to get depth and camera poses
     global vggt4track_model
     global tracker_model
     global wan_pipeline
@@ -234,17 +366,14 @@ def run_spatial_tracker(video_tensor: torch.Tensor):
     video_tensor_gpu = video_input.squeeze()
     unc_metric = depth_conf.squeeze().cpu().numpy() > 0.5
 
-    # Setup tracker
     tracker_model.spatrack.track_num = 512
     tracker_model.to("cuda")
 
-    # Get grid points for tracking
     frame_H, frame_W = video_tensor_gpu.shape[2:]
     grid_pts = get_points_on_a_grid(30, (frame_H, frame_W), device="cpu")
     query_xyt = torch.cat([torch.zeros_like(grid_pts[:, :, :1]), grid_pts], dim=2)[
         0].numpy()
 
-    # Run tracker
     with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
         (
             c2w_traj, intrs_out, point_map, conf_depth,
@@ -259,7 +388,6 @@ def run_spatial_tracker(video_tensor: torch.Tensor):
             support_frame=len(video_tensor_gpu)-1, replace_ratio=0.2
         )
 
-    # Resize outputs for rendering
     max_size = 384
     h, w = video_out.shape[2:]
     scale = min(max_size / h, max_size / w)
@@ -269,14 +397,19 @@ def run_spatial_tracker(video_tensor: torch.Tensor):
         point_map = T.Resize((new_h, new_w))(point_map)
         conf_depth = T.Resize((new_h, new_w))(conf_depth)
         intrs_out[:, :2, :] = intrs_out[:, :2, :] * scale
+        # Scale 2D track coordinates
+        track2d_pred[..., :2] = track2d_pred[..., :2] * scale
 
-    # Move results to CPU and return
     return {
         'video_out': video_out.cpu(),
         'point_map': point_map.cpu(),
         'conf_depth': conf_depth.cpu(),
         'intrs_out': intrs_out.cpu(),
         'c2w_traj': c2w_traj.cpu(),
+        'track2d_pred': track2d_pred.cpu(),  # 2D track positions (T, N, 2)
+        'track3d_pred': track3d_pred.cpu(),  # 3D track positions (T, N, 3)
+        'vis_pred': vis_pred.cpu(),          # Visibility mask (T, N)
+        'conf_pred': conf_pred.cpu(),        # Confidence scores (T, N)
     }
 
 
@@ -301,7 +434,6 @@ def run_wan_ttm_generation(prompt, tweak_index, tstrong_index, first_frame_path,
     progress(0.2, desc="Preparing inputs...")
     image = load_image(first_frame_path)
 
-    # Standard Wan Negative Prompt
     negative_prompt = (
         "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，"
         "低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，"
@@ -310,7 +442,6 @@ def run_wan_ttm_generation(prompt, tweak_index, tstrong_index, first_frame_path,
 
     wan_pipeline.to("cuda")
 
-    # Match resolution logic from run_wan.py
     max_area = 480 * 832
     mod_value = wan_pipeline.vae_scale_factor_spatial * \
         wan_pipeline.transformer.config.patch_size[1]
@@ -349,7 +480,7 @@ def run_wan_ttm_generation(prompt, tweak_index, tstrong_index, first_frame_path,
 
 def process_video(video_path, camera_movement, generate_ttm=True, progress=gr.Progress()):
     if video_path is None:
-        return None, None, None, None, "❌ Please upload a video first"
+        return None, None, None, None, None, "❌ Please upload a video first"
 
     progress(0, desc="Initializing...")
     temp_dir = create_user_temp_dir()
@@ -383,6 +514,33 @@ def process_video(video_path, camera_movement, generate_ttm=True, progress=gr.Pr
         new_exts = generate_camera_trajectory(len(
             rgb_frames), camera_movement, tracking_results['intrs_out'].numpy(), scene_scale)
 
+        progress(0.7, desc="Visualizing tracks...")
+        # Get track data for visualization
+        track2d = tracking_results['track2d_pred'].numpy()  # (T, N, 2+)
+        track3d = tracking_results['track3d_pred'].numpy()  # (T, N, 3)
+        vis_pred = tracking_results['vis_pred'].numpy()     # (T, N)
+
+        # Get depth at each track point (use Z coordinate from 3D tracks)
+        track_depths = track3d[..., 2]  # (T, N) - depth is the Z coordinate
+
+        # Create track visualization with selective history:
+        # - Background points (farther): show history trails
+        # - Foreground points (closer): no history trails
+        track_viz_path = os.path.join(out_dir, "track_visualization.mp4")
+        visualize_tracks_with_selective_history(
+            video=rgb_frames.copy(),
+            tracks=track2d[..., :2],  # Use only x, y coordinates
+            track_depths=track_depths,
+            visibility=vis_pred,
+            output_path=track_viz_path,
+            fps=OUTPUT_FPS,
+            depth_threshold=None,  # Auto-compute based on median depth
+            bg_history_length=-1,  # Infinite history for background
+            fg_history_length=0,   # No history for foreground
+            linewidth=2,
+            point_radius=4
+        )
+
         progress(0.8, desc="Rendering viewpoint...")
         output_video_path = os.path.join(out_dir, "rendered_video.mp4")
         render_results = render_from_pointcloud(rgb_frames, depth_frames, tracking_results['intrs_out'].numpy(),
@@ -395,11 +553,13 @@ def process_video(video_path, camera_movement, generate_ttm=True, progress=gr.Pr
             rgb_frames[0], cv2.COLOR_RGB2BGR))
 
         status_msg = f"✅ 3D results ready! You can now use the prompt below to generate a high-quality TTM video."
-        return render_results['rendered'], render_results['motion_signal'], render_results['mask'], first_frame_path, status_msg
+        return render_results['rendered'], render_results['motion_signal'], render_results['mask'], first_frame_path, track_viz_path, status_msg
 
     except Exception as e:
         logger.error(f"Error: {e}")
-        return None, None, None, None, f"❌ Error: {str(e)}"
+        import traceback
+        traceback.print_exc()
+        return None, None, None, None, None, f"❌ Error: {str(e)}"
 
 
 # --- GRADIO INTERFACE ---
@@ -408,7 +568,6 @@ with gr.Blocks(theme=gr.themes.Soft(), title="🎬 TTM Wan Video Generator") as
     gr.Markdown(
         "Transform standard videos into 3D-aware motion signals for Time-to-Move (TTM) generation.")
 
-    # Shared state for TTM files - initialized as empty strings
     first_frame_file = gr.State("")
     motion_signal_file = gr.State("")
     mask_file = gr.State("")
@@ -426,6 +585,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="🎬 TTM Wan Video Generator") as
                 "🚀 1. Run Spatial Tracker", variant="primary")
 
             output_video = gr.Video(label="Point Cloud Render (Draft)")
+            track_viz_output = gr.Video(label="Track Visualization (BG history, FG no history)")
             status_text = gr.Markdown("Ready...")
 
         with gr.Column(scale=1):
@@ -446,19 +606,14 @@ with gr.Blocks(theme=gr.themes.Soft(), title="🎬 TTM Wan Video Generator") as
             wan_output_video = gr.Video(label="Final High-Quality TTM Video")
             wan_status = gr.Markdown("Awaiting 3D inputs...")
 
-    # The Accordion provides a visual check of what TTM is using
     with gr.Accordion("Debug: TTM Intermediate Inputs", open=False):
         with gr.Row():
-            # IMPORTANT: type="filepath" prevents the ValueError by passing
-            # the path string instead of the raw pixel array.
             motion_signal_output = gr.Video(label="motion_signal.mp4")
             mask_output = gr.Video(label="mask.mp4")
             first_frame_output = gr.Image(
                 label="first_frame.png", type="filepath")
 
-    # --- Event Handlers ---
 
-    # 1. Process 3D Tracking and save results to temporary local files
     generate_btn.click(
         fn=process_video,
         inputs=[video_input, camera_movement],
@@ -467,23 +622,22 @@ with gr.Blocks(theme=gr.themes.Soft(), title="🎬 TTM Wan Video Generator") as
             motion_signal_output,
             mask_output,
             first_frame_output,
+            track_viz_output,
             status_text
         ]
     ).then(
-        # 2. Update the State variables with the file paths from the previous step.
-        # We ignore the 'output_video' (index 0) and 'status_text' (index 4).
-        fn=lambda a, b, c, d, e: (b, c, d),
+        fn=lambda a, b, c, d, e, f: (b, c, d),
         inputs=[
             output_video,
             motion_signal_output,
             mask_output,
             first_frame_output,
+            track_viz_output,
             status_text
         ],
         outputs=[motion_signal_file, mask_file, first_frame_file]
     )
 
-    # 3. Use the stored paths to run the Wan 2.2 TTM Dual-Clock Denoising loop
     wan_generate_btn.click(
         fn=run_wan_ttm_generation,
         inputs=[