Update visualizer.py

visualizer.py

@@ -104,7 +104,7 @@ class Visualizer:
     Handles drawing tracked points and their trajectories on video frames.
 
     Args:
-        output_dir: Directory to save output visualizations
+        save_dir: Directory to save output visualizations
         padding: Padding to add around video frames in pixels
         fps: Frames per second for output video
         colormap: Color scheme for tracks ('rainbow' or 'spring')
@@ -115,20 +115,21 @@ class Visualizer:
 
     def __init__(
         self,
-        output_dir: str = "./results",
-        padding: int = 0,
+        save_dir: str = "./results",
+        grayscale: bool = False,
+        pad_value: int = 0,
         fps: int = 10,
         colormap: str = "rainbow",
-        line_width: int = 2,
-        initial_frame_repeat: int = 10,
-        track_history_length: int = 0,
+        linewidth: int = 2,
+        show_first_frame: int = 10,
+        tracks_leave_trace: int = 0,
     ):
-        self.output_dir = output_dir
-        self.padding = padding
+        self.save_dir = save_dir
+        self.padding = pad_value
         self.fps = fps
-        self.line_width = line_width
-        self.initial_frame_repeat = initial_frame_repeat
-        self.track_history_length = track_history_length
+        self.line_width = linewidth
+        self.initial_frame_repeat = show_first_frame
+        self.track_history_length = tracks_leave_trace
 
         # Set up colormap for track visualization
         self.colormap = colormap
@@ -143,6 +144,7 @@ class Visualizer:
         self,
         video: torch.Tensor,  # (B,T,C,H,W)
         tracks: torch.Tensor,  # (B,T,N,2)
+        
         visibility: torch.Tensor = None,  # (B,T,N,1) bool
         segmentation: torch.Tensor = None,  # (B,1,H,W)
         filename: str = "video",
@@ -209,7 +211,7 @@ class Visualizer:
             video: Video tensor of shape (B,T,C,H,W)
             filename: Output filename without extension
         """
-        os.makedirs(self.output_dir, exist_ok=True)
+        os.makedirs(self.save_dir, exist_ok=True)
 
         # Extract frames from video tensor
         frames = [
@@ -217,7 +219,7 @@ class Visualizer:
             for frame in video.unbind(1)
         ]
 
-        output_path = os.path.join(self.output_dir, f"{filename}.mp4")
+        output_path = os.path.join(self.save_dir, f"{filename}.mp4")
 
         try:
             with imageio.get_writer(output_path, fps=self.fps, quality=8) as writer: