revert

app.py

@@ -24,7 +24,7 @@ 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
+
 
 from diffusers.utils import export_to_video, load_image
 
@@ -194,145 +194,6 @@ 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):
     """
@@ -397,8 +258,6 @@ 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
 
     return {
         'video_out': video_out.cpu(),
@@ -406,10 +265,6 @@ def run_spatial_tracker(video_tensor: torch.Tensor):
         '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)
     }
 
 
@@ -480,7 +335,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, None, "❌ Please upload a video first"
+        return None, None, None, None, "❌ Please upload a video first"
 
     progress(0, desc="Initializing...")
     temp_dir = create_user_temp_dir()
@@ -514,33 +369,6 @@ 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(),
@@ -553,13 +381,11 @@ 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, track_viz_path, status_msg
+        return render_results['rendered'], render_results['motion_signal'], render_results['mask'], first_frame_path, status_msg
 
     except Exception as e:
         logger.error(f"Error: {e}")
-        import traceback
-        traceback.print_exc()
-        return None, None, None, None, None, f"❌ Error: {str(e)}"
+        return None, None, None, None, f"❌ Error: {str(e)}"
 
 
 # --- GRADIO INTERFACE ---
@@ -585,7 +411,6 @@ 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):
@@ -622,17 +447,15 @@ 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(
-        fn=lambda a, b, c, d, e, f: (b, c, d),
+        fn=lambda a, b, c, d, e: (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]