Update visualizer.py

visualizer.py

@@ -4,555 +4,390 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
-
 import os
-from typing import List
-
-import imageio
-import matplotlib.pyplot as plt
 import numpy as np
+import imageio
 import torch
-import torch.nn.functional as F
 
 from matplotlib import cm
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+import matplotlib.pyplot as plt
 from PIL import Image, ImageDraw
 
 
-def draw_circle_on_image(
-    image: Image,
-    center: tuple,
-    radius: int,
-    color: tuple = (255, 0, 0),
-    visible: bool = True,
-    alpha: int = None,
-) -> Image:
-    """Draw a circle on a PIL Image.
-
-    Args:
-        image: PIL Image to draw on
-        center: (x,y) coordinates of circle center
-        radius: Radius of circle in pixels
-        color: RGB color tuple
-        visible: Whether to fill the circle
-        alpha: Optional alpha value for transparency
-
-    Returns:
-        Modified PIL Image
-    """
-    draw = ImageDraw.Draw(image, 'RGBA')  # Enable alpha channel
-    # Use float coordinates for smoother rendering
-    bbox = [
-        (center[0] - radius, center[1] - radius),
-        (center[0] + radius, center[1] + radius),
-    ]
-    color = tuple(list(color) + [alpha if alpha is not None else 255])
-    
-    # Use anti-aliasing by drawing a slightly larger circle underneath
-    if visible:
-        # Draw a slightly larger background circle for anti-aliasing
-        larger_bbox = [
-            (center[0] - radius - 0.5, center[1] - radius - 0.5),
-            (center[0] + radius + 0.5, center[1] + radius + 0.5),
-        ]
-        draw.ellipse(larger_bbox, fill=tuple(list(color[:-1]) + [int(color[-1] * 0.5)]))
-    
-    draw.ellipse(bbox, fill=tuple(color) if visible else None, outline=tuple(color))
-    return image
-
-
-def draw_line_segment(
-    image: Image, start: tuple, end: tuple, color: tuple, width: int
-) -> Image:
-    """Draw a line on a PIL Image.
-
-    Args:
-        image: PIL Image to draw on
-        start: (x,y) coordinates of line start
-        end: (x,y) coordinates of line end
-        color: RGB color tuple
-        width: Line width in pixels
-
-    Returns:
-        Modified PIL Image
-    """
-    draw = ImageDraw.Draw(image)
-    draw.line((start[0], start[1], end[0], end[1]), fill=tuple(color), width=width)
-    return image
-
-
-def blend_images(
-    image1: np.ndarray, alpha: float, image2: np.ndarray, beta: float, gamma: float
-) -> np.ndarray:
-    """Blend two images with weights.
-
-    Args:
-        image1: First image array
-        alpha: Weight of first image
-        image2: Second image array
-        beta: Weight of second image
-        gamma: Scalar added to weighted sum
-
-    Returns:
-        Blended uint8 image array
-    """
-    return (image1 * alpha + image2 * beta + gamma).astype("uint8")
+def read_video_from_path(path):
+    try:
+        reader = imageio.get_reader(path)
+    except Exception as e:
+        print("Error opening video file: ", e)
+        return None
+    frames = []
+    for i, im in enumerate(reader):
+        frames.append(np.array(im))
+    return np.stack(frames)
 
 
-class Visualizer:
-    """A class for visualizing point tracks on videos.
+def draw_circle(rgb, coord, radius, color=(255, 0, 0), visible=True, color_alpha=None):
+    # Create a draw object
+    draw = ImageDraw.Draw(rgb)
+    # Calculate the bounding box of the circle
+    left_up_point = (coord[0] - radius, coord[1] - radius)
+    right_down_point = (coord[0] + radius, coord[1] + radius)
+    # Draw the circle
+    color = tuple(list(color) + [color_alpha if color_alpha is not None else 255])
+
+    draw.ellipse(
+        [left_up_point, right_down_point],
+        fill=tuple(color) if visible else None,
+        outline=tuple(color),
+    )
+    return rgb
 
-    Handles drawing tracked points and their trajectories on video frames.
 
-    Args:
-        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')
-        line_width: Width of track lines in pixels
-        initial_frame_repeat: Number of times to repeat first frame
-        track_history_length: How many past frames to show tracks for (0=current only, -1=all)
-    """
+def draw_line(rgb, coord_y, coord_x, color, linewidth):
+    draw = ImageDraw.Draw(rgb)
+    draw.line(
+        (coord_y[0], coord_y[1], coord_x[0], coord_x[1]),
+        fill=tuple(color),
+        width=linewidth,
+    )
+    return rgb
 
+
+def add_weighted(rgb, alpha, original, beta, gamma):
+    return (rgb * alpha + original * beta + gamma).astype("uint8")
+
+
+class Visualizer:
     def __init__(
         self,
         save_dir: str = "./results",
         grayscale: bool = False,
         pad_value: int = 0,
         fps: int = 10,
-        colormap: str = "rainbow",
+        mode: str = "rainbow",  # 'cool', 'optical_flow'
         linewidth: int = 2,
         show_first_frame: int = 10,
-        tracks_leave_trace: int = 0,
+        tracks_leave_trace: int = 0,  # -1 for infinite
     ):
+        self.mode = mode
         self.save_dir = save_dir
-        self.padding = pad_value
+        if mode == "rainbow":
+            self.color_map = cm.get_cmap("gist_rainbow")
+        elif mode == "cool":
+            self.color_map = cm.get_cmap(mode)
+        self.show_first_frame = show_first_frame
+        self.grayscale = grayscale
+        self.tracks_leave_trace = tracks_leave_trace
+        self.pad_value = pad_value
+        self.linewidth = linewidth
         self.fps = fps
-        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
-        if colormap not in ["rainbow", "spring"]:
-            raise ValueError("Colormap must be 'rainbow' or 'spring'")
-
-        self.color_mapper = cm.get_cmap(
-            "gist_rainbow" if colormap == "rainbow" else "spring"
-        )
 
     def visualize(
         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)
+        visibility: torch.Tensor = None,  # (B, T, N, 1) bool
+        gt_tracks: torch.Tensor = None,  # (B,T,N,2)
+        segm_mask: torch.Tensor = None,  # (B,1,H,W)
         filename: str = "video",
-        query_frame: int = 0,
+        writer=None,  # tensorboard Summary Writer, used for visualization during training
+        step: int = 0,
+        query_frame=0,
         save_video: bool = True,
-        point_opacity: float = 1.0,
-    ) -> torch.Tensor:
-        """Visualize tracked points and their trajectories on video frames.
-
-        Args:
-            video: Input video tensor of shape (B,T,C,H,W)
-            tracks: Point track coordinates of shape (B,T,N,2)
-            visibility: Optional visibility mask of shape (B,T,N,1)
-            segmentation: Optional segmentation mask of shape (B,1,H,W)
-            filename: Output filename for saved video
-            query_frame: Frame index to use for color assignment
-            save_video: Whether to save visualization video
-            point_opacity: Opacity value for track points (0-1)
-
-        Returns:
-            Tensor containing visualization frames
-        """
-        # Process segmentation if provided
-        if segmentation is not None:
-            coords = tracks[0, query_frame].round().long()
-            segmentation = segmentation[0, query_frame][
-                coords[:, 1], coords[:, 0]
-            ].long()
-
-        # Add padding to video frames
-        padded_video = F.pad(
+        compensate_for_camera_motion: bool = False,
+        opacity: float = 1.0,
+    ):
+        if compensate_for_camera_motion:
+            assert segm_mask is not None
+        # if segm_mask is not None:
+        #     coords = tracks[0, query_frame].round().long()
+        #     segm_mask = segm_mask[0, query_frame][coords[:, 1], coords[:, 0]].long()
+
+        video = F.pad(
             video,
-            (self.padding, self.padding, self.padding, self.padding),
-            mode="constant",
-            value=255,
+            (self.pad_value, self.pad_value, self.pad_value, self.pad_value),
+            "constant",
+            255,
         )
+        color_alpha = int(opacity * 255)
+        tracks = tracks + self.pad_value
 
-        # Convert opacity to integer value
-        opacity_value = min(max(int(point_opacity * 255), 0), 255)
+        if self.grayscale:
+            transform = transforms.Grayscale()
+            video = transform(video)
+            video = video.repeat(1, 1, 3, 1, 1)
 
-        # Adjust track coordinates for padding
-        padded_tracks = tracks + self.padding
-
-        # Generate visualization frames
-        output_video = self.draw_tracks_on_video(
-            video=padded_video,
-            tracks=padded_tracks,
+        res_video = self.draw_tracks_on_video(
+            video=video,
+            tracks=tracks,
             visibility=visibility,
-            segmentation=segmentation,
+            segm_mask=segm_mask,
+            gt_tracks=gt_tracks,
             query_frame=query_frame,
-            opacity=opacity_value,
+            compensate_for_camera_motion=compensate_for_camera_motion,
+            color_alpha=color_alpha,
         )
-
-        # Save video if requested
         if save_video:
-            self.save_video(output_video, filename=filename)
-
-        return output_video
-
-    def save_video(self, video: torch.Tensor, filename: str):
-        """Save video tensor as MP4 file.
-
-        Args:
-            video: Video tensor of shape (B,T,C,H,W)
-            filename: Output filename without extension
-        """
-        os.makedirs(self.save_dir, exist_ok=True)
+            self.save_video(res_video, filename=filename, writer=writer, step=step)
+        return res_video
+
+    def save_video(self, video, filename, writer=None, step=0):
+        if writer is not None:
+            writer.add_video(
+                filename,
+                video.to(torch.uint8),
+                global_step=step,
+                fps=self.fps,
+            )
+        else:
+            os.makedirs(self.save_dir, exist_ok=True)
+            wide_list = list(video.unbind(1))
+            wide_list = [wide[0].permute(1, 2, 0).cpu().numpy() for wide in wide_list]
 
-        # Extract frames from video tensor
-        frames = [
-            frame[0].permute(1, 2, 0).cpu().numpy().astype(np.uint8)
-            for frame in video.unbind(1)
-        ]
+            # Prepare the video file path
+            save_path = os.path.join(self.save_dir, f"{filename}.mp4")
 
-        output_path = os.path.join(self.save_dir, f"{filename}.mp4")
+            # Create a writer object
+            video_writer = imageio.get_writer(save_path, fps=self.fps)
 
-        try:
-            with imageio.get_writer(output_path, fps=self.fps, quality=8) as writer:
-                # Write frames excluding padding frames
-                for frame in frames[2:-1]:
-                    writer.append_data(frame)
+            # Write frames to the video file
+            for frame in wide_list[2:-1]:
+                video_writer.append_data(frame)
 
-            print(f"Successfully saved video to {output_path}")
+            video_writer.close()
 
-        except Exception as e:
-            print(f"Error saving video to {output_path}: {str(e)}")
+            print(f"Video saved to {save_path}")
 
     def draw_tracks_on_video(
         self,
         video: torch.Tensor,
         tracks: torch.Tensor,
         visibility: torch.Tensor = None,
-        segmentation: torch.Tensor = None,
-        query_frame: int = 0,
-        opacity: int = 255,
-    ) -> torch.Tensor:
-        """Draw tracks on video frames.
-
-        Args:
-            video: Video tensor of shape (B,T,C,H,W)
-            tracks: Track coordinates tensor of shape (B,T,N,2)
-            visibility: Optional visibility mask of shape (B,T,N)
-            segmentation: Optional segmentation mask for coloring
-            query_frame: Frame index to use for rainbow coloring
-            opacity: Opacity value for track points (0-255)
-
-        Returns:
-            Video tensor with visualized tracks of shape (1,T,3,H,W)
-        """
-        # Validate input dimensions
-        _, num_frames, channels, _, _ = video.shape
-        _, _, num_points, dims = tracks.shape
-        assert dims == 2 and channels == 3, "Invalid input dimensions"
-
-        # Convert tensors to numpy arrays but keep as float
-        video_np = video[0].permute(0, 2, 3, 1).byte().detach().cpu().numpy()
-        tracks_np = tracks[0].detach().cpu().numpy()
-
-        # Create output frame buffer
-        output_frames = [frame.copy() for frame in video_np]
-
-        # Assign colors to tracks based on segmentation or position
-        track_colors = (
-            self._assign_segmentation_colors(
-                tracks_np, segmentation, num_frames, num_points
-            )
-            if segmentation is not None
-            else self._assign_track_colors(
-                tracks_np, query_frame, num_frames, num_points
-            )
-        )
-
-        # Draw track history lines if enabled
-        if self.track_history_length != 0:
-            output_frames = self._draw_track_lines(
-                output_frames, tracks_np, track_colors, query_frame, num_frames
-            )
-
-        # Draw track points with visibility and opacity
-        output_frames = self._draw_track_points(
-            output_frames, tracks_np, track_colors, visibility, opacity
-        )
-
-        # Add initial frame repeats for better visualization
-        if self.initial_frame_repeat > 0:
-            output_frames = [
-                output_frames[0]
-            ] * self.initial_frame_repeat + output_frames[1:]
-
-        # Convert back to torch tensor
-        return (
-            torch.from_numpy(np.stack(output_frames)).permute(0, 3, 1, 2)[None].byte()
-        )
-
-    def _assign_track_colors(
-        self, tracks: np.ndarray, query_frame: int, num_frames: int, num_points: int
-    ) -> np.ndarray:
-        """Assigns colors to tracks based on either rainbow mapping of y-coordinates or time-based coloring.
-
-        Args:
-            tracks: Track coordinates array of shape (num_frames, num_points, 2)
-            query_frame: Frame index to use for rainbow coloring
-            num_frames: Total number of frames
-            num_points: Number of tracked points
-
-        Returns:
-            Array of track colors with shape (num_frames, num_points, 3)
-        """
-        track_colors = np.zeros((num_frames, num_points, 3))
-
-        if self.colormap == "rainbow":
-            # Normalize y-coordinates to [0,1] range for rainbow coloring
-            y_coords = tracks[query_frame, :, 1]
-            y_min, y_max = y_coords.min(), y_coords.max()
-            if y_min == y_max:
-                y_max = y_min + 1  # Avoid division by zero
-            norm = plt.Normalize(y_min, y_max)
-
-            # Assign colors based on normalized y-coordinate
-            for point_idx in range(num_points):
-                query_idx = (
-                    query_frame[point_idx]
-                    if isinstance(query_frame, torch.Tensor)
-                    else query_frame
-                )
-                color = (
-                    np.array(
-                        self.color_mapper(norm(tracks[query_idx, point_idx, 1]))[:3]
-                    )[None]
-                    * 255
-                )
-                track_colors[:, point_idx] = np.repeat(color, num_frames, axis=0)
-        else:
-            # Assign colors that vary smoothly with time
-            for frame_idx in range(num_frames):
-                color = (
-                    np.array(self.color_mapper(frame_idx / max(1, num_frames - 1))[:3])[
-                        None
-                    ]
-                    * 255
+        segm_mask: torch.Tensor = None,
+        gt_tracks=None,
+        query_frame=0,
+        compensate_for_camera_motion=False,
+        color_alpha: int = 255,
+    ):
+        B, T, C, H, W = video.shape
+        _, _, N, D = tracks.shape
+
+        assert D == 2
+        assert C == 3
+        video = video[0].permute(0, 2, 3, 1).byte().detach().cpu().numpy()  # S, H, W, C
+        tracks = tracks[0].detach().cpu().numpy()  # S, N, 2
+        if gt_tracks is not None:
+            gt_tracks = gt_tracks[0].detach().cpu().numpy()
+
+        res_video = []
+
+        # process input video
+        for rgb in video:
+            res_video.append(rgb.copy())
+        vector_colors = np.zeros((T, N, 3))
+
+        if self.mode == "optical_flow":
+            import flow_vis
+
+            vector_colors = flow_vis.flow_to_color(tracks - tracks[query_frame][None])
+        elif segm_mask is None:
+            if self.mode == "rainbow":
+                y_min, y_max = (
+                    tracks[query_frame, :, 1].min(),
+                    tracks[query_frame, :, 1].max(),
                 )
-                track_colors[frame_idx] = np.repeat(color, num_points, axis=0)
-
-        return track_colors.astype(np.uint8)
-
-    def _assign_segmentation_colors(
-        self,
-        tracks: np.ndarray,
-        segmentation: torch.Tensor,
-        num_frames: int,
-        num_points: int,
-    ) -> np.ndarray:
-        """Assigns colors to tracks based on segmentation masks and colormap.
-
-        Args:
-            tracks: Track coordinates array of shape (num_frames, num_points, 2)
-            segmentation: Binary segmentation mask of shape (num_points,)
-            num_frames: Total number of frames
-            num_points: Number of tracked points
-
-        Returns:
-            Array of track colors with shape (num_frames, num_points, 3)
-        """
-        track_colors = np.zeros((num_frames, num_points, 3))
-
-        if self.colormap == "rainbow":
-            # Set background points to white
-            background_mask = segmentation <= 0
-            track_colors[:, background_mask, :] = 255
-
-            # Color foreground points based on y-coordinate
-            foreground_mask = segmentation > 0
-            if torch.any(foreground_mask):
-                y_coords = tracks[0, foreground_mask, 1]
-                y_min, y_max = y_coords.min(), y_coords.max()
-                if y_min == y_max:
-                    y_max = y_min + 1  # Avoid division by zero
                 norm = plt.Normalize(y_min, y_max)
-
-                for point_idx in range(num_points):
-                    if segmentation[point_idx] > 0:
-                        color = (
-                            np.array(
-                                self.color_mapper(norm(tracks[0, point_idx, 1]))[:3]
-                            )[None]
-                            * 255
-                        )
-                        track_colors[:, point_idx] = np.repeat(
-                            color, num_frames, axis=0
-                        )
+                for n in range(N):
+                    if isinstance(query_frame, torch.Tensor):
+                        query_frame_ = query_frame[n]
+                    else:
+                        query_frame_ = query_frame
+                    color = self.color_map(norm(tracks[query_frame_, n, 1]))
+                    color = np.array(color[:3])[None] * 255
+                    vector_colors[:, n] = np.repeat(color, T, axis=0)
+            else:
+                # color changes with time
+                for t in range(T):
+                    color = np.array(self.color_map(t / T)[:3])[None] * 255
+                    vector_colors[t] = np.repeat(color, N, axis=0)
         else:
-            # Binary coloring based on segmentation
-            segmentation = segmentation.cpu()
-            colors = np.zeros((num_points, 3), dtype=np.float32)
-            colors[segmentation > 0] = (
-                np.array(self.color_mapper(1.0)[:3]) * 255.0
-            )  # Foreground
-            colors[segmentation <= 0] = (
-                np.array(self.color_mapper(0.0)[:3]) * 255.0
-            )  # Background
-            track_colors = np.repeat(colors[None], num_frames, axis=0)
-
-        return track_colors.astype(np.uint8)
-
-    def _draw_track_lines(
-        self,
-        frames: List[np.ndarray],
-        tracks: np.ndarray,
-        track_colors: np.ndarray,
-        query_frame: int,
-        num_frames: int,
-    ) -> List[np.ndarray]:
-        """Draw track lines showing point trajectories over time.
-
-        Args:
-            frames: List of video frames to draw on
-            tracks: Array of track coordinates (num_frames, num_points, 2)
-            track_colors: Array of track colors (num_frames, num_points, 3)
-            query_frame: Frame index where tracking starts
-            num_frames: Total number of frames
-
-        Returns:
-            List of frames with track lines drawn
-        """
-        # Draw tracks starting from query frame
-        for frame_idx in range(query_frame + 1, num_frames):
-            # Get track history based on history length setting
-            start_idx = (
-                max(0, frame_idx - self.track_history_length)
-                if self.track_history_length >= 0
-                else 0
-            )
-
-            # Extract relevant track segments and colors
-            curr_tracks = tracks[start_idx : frame_idx + 1]
-            curr_colors = track_colors[start_idx : frame_idx + 1]
-
-            # Draw track segments on current frame
-            frames[frame_idx] = self._draw_track_segments(
-                frames[frame_idx], curr_tracks, curr_colors
-            )
+            if self.mode == "rainbow":
+                vector_colors[:, segm_mask <= 0, :] = 255
 
-        return frames
-
-    def _draw_track_segments(
-        self,
-        frame: np.ndarray,
-        tracks: np.ndarray,
-        colors: np.ndarray,
-    ) -> np.ndarray:
-        """Draw track segments showing point trajectories between consecutive frames.
-
-        Args:
-            frame: Video frame to draw on
-            tracks: Array of track coordinates (num_segments, num_points, 2)
-            colors: Array of track colors (num_segments, num_points, 3)
-
-        Returns:
-            Frame with track segments drawn
-        """
-        num_segments, num_points, _ = tracks.shape
-        frame_img = Image.fromarray(np.uint8(frame))
-
-        for segment_idx in range(num_segments - 1):
-            segment_color = colors[segment_idx]
-            original = frame_img.copy()
-
-            # Use cubic falloff for track history opacity
-            alpha = (segment_idx / num_segments) ** 3
-
-            valid_points = ~np.isclose(tracks[segment_idx], 0).all(axis=1)
-
-            for point_idx in range(num_points):
-                if valid_points[point_idx]:
-                    start = (
-                        tracks[segment_idx, point_idx, 0],
-                        tracks[segment_idx, point_idx, 1],
-                    )
-                    end = (
-                        tracks[segment_idx + 1, point_idx, 0],
-                        tracks[segment_idx + 1, point_idx, 1],
+                y_min, y_max = (
+                    tracks[0, segm_mask > 0, 1].min(),
+                    tracks[0, segm_mask > 0, 1].max(),
+                )
+                norm = plt.Normalize(y_min, y_max)
+                for n in range(N):
+                    if segm_mask[n] > 0:
+                        color = self.color_map(norm(tracks[0, n, 1]))
+                        color = np.array(color[:3])[None] * 255
+                        vector_colors[:, n] = np.repeat(color, T, axis=0)
+
+            else:
+                # color changes with segm class
+                segm_mask = segm_mask.cpu()
+                color = np.zeros((segm_mask.shape[0], 3), dtype=np.float32)
+                color[segm_mask > 0] = np.array(self.color_map(1.0)[:3]) * 255.0
+                color[segm_mask <= 0] = np.array(self.color_map(0.0)[:3]) * 255.0
+                vector_colors = np.repeat(color[None], T, axis=0)
+
+        #  draw tracks
+        if self.tracks_leave_trace != 0:
+            for t in range(query_frame + 1, T):
+                first_ind = (
+                    max(0, t - self.tracks_leave_trace)
+                    if self.tracks_leave_trace >= 0
+                    else 0
+                )
+                curr_tracks = tracks[first_ind : t + 1]
+                curr_colors = vector_colors[first_ind : t + 1]
+                if compensate_for_camera_motion:
+                    diff = (
+                        tracks[first_ind : t + 1, segm_mask <= 0]
+                        - tracks[t : t + 1, segm_mask <= 0]
+                    ).mean(1)[:, None]
+
+                    curr_tracks = curr_tracks - diff
+                    curr_tracks = curr_tracks[:, segm_mask > 0]
+                    curr_colors = curr_colors[:, segm_mask > 0]
+
+                res_video[t] = self._draw_pred_tracks(
+                    res_video[t],
+                    curr_tracks,
+                    curr_colors,
+                )
+                if gt_tracks is not None:
+                    res_video[t] = self._draw_gt_tracks(
+                        res_video[t], gt_tracks[first_ind : t + 1]
                     )
 
-                    frame_img = draw_line_segment(
-                        frame_img,
-                        start,
-                        end,
-                        segment_color[point_idx].astype(int),
-                        self.line_width,
+        #  draw points
+        for t in range(T):
+            img = Image.fromarray(np.uint8(res_video[t]))
+            for i in range(N):
+                coord = (tracks[t, i, 0], tracks[t, i, 1])
+                visibile = True
+                if visibility is not None:
+                    visibile = visibility[0, t, i]
+                if coord[0] != 0 and coord[1] != 0:
+                    if not compensate_for_camera_motion or (
+                        compensate_for_camera_motion and segm_mask[i] > 0
+                    ):
+                        # img = draw_circle(
+                        #     img,
+                        #     coord=coord,
+                        #     radius=int(self.linewidth * 2),
+                        #     color=vector_colors[t, i].astype(int),
+                        #     visible=visibile,
+                        #     color_alpha=color_alpha,
+                        # )
+
+                        # coord_ = coord[t,i]
+                        #  draw a red cross
+                        # if gt_tracks[0] > 0 and gt_tracks[1] > 0:
+                        if visibile:
+                            length = self.linewidth * 3
+                            coord_y = ((coord[0]) + length, (coord[1]) + length)
+                            coord_x = ((coord[0]) - length, (coord[1]) - length)
+                            rgb = draw_line(
+                                img,
+                                coord_y,
+                                coord_x,
+                                vector_colors[t, i].astype(int),
+                                self.linewidth,
+                            )
+                            coord_y = ((coord[0]) - length, (coord[1]) + length)
+                            coord_x = ((coord[0]) + length, (coord[1]) - length)
+                            rgb = draw_line(
+                                img,
+                                coord_y,
+                                coord_x,
+                                vector_colors[t, i].astype(int),
+                                self.linewidth,
+                            )
+            res_video[t] = np.array(img)
+
+        #  construct the final rgb sequence
+        if self.show_first_frame > 0:
+            res_video = [res_video[0]] * self.show_first_frame + res_video[1:]
+        return torch.from_numpy(np.stack(res_video)).permute(0, 3, 1, 2)[None].byte()
+
+    def _draw_pred_tracks(
+        self,
+        rgb: np.ndarray,  # H x W x 3
+        tracks: np.ndarray,  # T x 2
+        vector_colors: np.ndarray,
+        alpha: float = 0.5,
+    ):
+        T, N, _ = tracks.shape
+        rgb = Image.fromarray(np.uint8(rgb))
+        for s in range(T - 1):
+            vector_color = vector_colors[s]
+            original = rgb.copy()
+            alpha = (s / T) ** 2
+            for i in range(N):
+                coord_y = ((tracks[s, i, 0]), (tracks[s, i, 1]))
+                coord_x = ((tracks[s + 1, i, 0]), (tracks[s + 1, i, 1]))
+                if coord_y[0] != 0 and coord_y[1] != 0:
+                    rgb = draw_line(
+                        rgb,
+                        coord_y,
+                        coord_x,
+                        vector_color[i].astype(int),
+                        self.linewidth,
                     )
-
-            if self.track_history_length > 0:
-                frame_img = Image.fromarray(
-                    blend_images(
-                        np.array(frame_img), alpha, np.array(original), 1 - alpha, 0
+            if self.tracks_leave_trace > 0:
+                rgb = Image.fromarray(
+                    np.uint8(
+                        add_weighted(
+                            np.array(rgb), alpha, np.array(original), 1 - alpha, 0
+                        )
                     )
                 )
+        rgb = np.array(rgb)
+        return rgb
 
-        return np.array(frame_img)
-
-    def _draw_track_points(
+    def _draw_gt_tracks(
         self,
-        frames: List[np.ndarray],
-        tracks: np.ndarray,
-        track_colors: np.ndarray,
-        visibility: torch.Tensor,
-        opacity: int,
-    ) -> List[np.ndarray]:
-        """Draw tracked points on each frame with circles.
-
-        Args:
-            frames: List of video frames to draw on
-            tracks: Array of track coordinates (num_frames, num_points, 2)
-            track_colors: Array of track colors (num_frames, num_points, 3)
-            visibility: Tensor indicating point visibility per frame
-            opacity: Opacity value for drawing points
-
-        Returns:
-            List of frames with track points drawn
-        """
-        frame_imgs = [Image.fromarray(np.uint8(frame)) for frame in frames]
-        
-        # Use more precise validation of points
-        valid_points = ~np.isclose(tracks, 0).all(axis=2)
-
-        for frame_idx, frame_img in enumerate(frame_imgs):
-            frame_visibility = (
-                np.ones(tracks.shape[1], dtype=bool)
-                if visibility is None
-                else visibility[0, frame_idx].cpu().numpy()
-            )
-
-            points_to_draw = np.logical_and(valid_points[frame_idx], frame_visibility)
-
-            for point_idx in np.where(points_to_draw)[0]:
-                # Keep coordinates as floats
-                coord = tuple(tracks[frame_idx, point_idx])
-                color = track_colors[frame_idx, point_idx].astype(int)
-
-                frame_img = draw_circle_on_image(
-                    frame_img,
-                    center=coord,
-                    radius=int(self.line_width * 2),
-                    color=color,
-                    visible=frame_visibility[point_idx],
-                    alpha=opacity,
-                )
-
-            frames[frame_idx] = np.array(frame_img)
-
-        return frames
+        rgb: np.ndarray,  # H x W x 3,
+        gt_tracks: np.ndarray,  # T x 2
+        vector_colors: np.ndarray = None,
+    ):
+        T, N, _ = gt_tracks.shape
+        if vector_colors is None:
+            color = np.array((211, 0, 0))
+        rgb = Image.fromarray(np.uint8(rgb))
+        for t in range(T):
+            if vector_colors is not None:
+                vector_color = vector_colors[t]
+            for i in range(N):
+                if vector_colors is not None:
+                    color = vector_color[i].astype(int)
+                gt_tracks = gt_tracks[t][i]
+                #  draw a red cross
+                if gt_tracks[0] > 0 and gt_tracks[1] > 0:
+                    length = self.linewidth * 3
+                    coord_y = ((gt_tracks[0]) + length, (gt_tracks[1]) + length)
+                    coord_x = ((gt_tracks[0]) - length, (gt_tracks[1]) - length)
+                    rgb = draw_line(
+                        rgb,
+                        coord_y,
+                        coord_x,
+                        color,
+                        self.linewidth,
+                    )
+                    coord_y = ((gt_tracks[0]) - length, (gt_tracks[1]) + length)
+                    coord_x = ((gt_tracks[0]) + length, (gt_tracks[1]) - length)
+                    rgb = draw_line(
+                        rgb,
+                        coord_y,
+                        coord_x,
+                        color,
+                        self.linewidth,
+                    )
+        rgb = np.array(rgb)
+        return rgb