Update visualizer.py

visualizer.py

@@ -4,390 +4,553 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
+
 import os
-import numpy as np
+from typing import List
+
 import imageio
+import matplotlib.pyplot as plt
+import numpy as np
 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 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)
-
-
-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])
+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")
 
-    draw.ellipse(
-        [left_up_point, right_down_point],
-        fill=tuple(color) if visible else None,
-        outline=tuple(color),
-    )
-    return rgb
-
-
-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
 
+class Visualizer:
+    """A class for visualizing point tracks on videos.
 
-def add_weighted(rgb, alpha, original, beta, gamma):
-    return (rgb * alpha + original * beta + gamma).astype("uint8")
+    Handles drawing tracked points and their trajectories on video frames.
 
+    Args:
+        output_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)
+    """
 
-class Visualizer:
     def __init__(
         self,
-        save_dir: str = "./results",
-        grayscale: bool = False,
-        pad_value: int = 0,
+        output_dir: str = "./results",
+        padding: int = 0,
         fps: int = 10,
-        mode: str = "rainbow",  # 'cool', 'optical_flow'
-        linewidth: int = 2,
-        show_first_frame: int = 10,
-        tracks_leave_trace: int = 0,  # -1 for infinite
+        colormap: str = "rainbow",
+        line_width: int = 2,
+        initial_frame_repeat: int = 10,
+        track_history_length: int = 0,
     ):
-        self.mode = mode
-        self.save_dir = save_dir
-        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.output_dir = output_dir
+        self.padding = padding
         self.fps = fps
+        self.line_width = line_width
+        self.initial_frame_repeat = initial_frame_repeat
+        self.track_history_length = track_history_length
+
+        # 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
-        gt_tracks: torch.Tensor = None,  # (B,T,N,2)
-        segm_mask: torch.Tensor = None,  # (B,1,H,W)
+        visibility: torch.Tensor = None,  # (B,T,N,1) bool
+        segmentation: torch.Tensor = None,  # (B,1,H,W)
         filename: str = "video",
-        writer=None,  # tensorboard Summary Writer, used for visualization during training
-        step: int = 0,
-        query_frame=0,
+        query_frame: int = 0,
         save_video: bool = True,
-        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(
+        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(
             video,
-            (self.pad_value, self.pad_value, self.pad_value, self.pad_value),
-            "constant",
-            255,
+            (self.padding, self.padding, self.padding, self.padding),
+            mode="constant",
+            value=255,
         )
-        color_alpha = int(opacity * 255)
-        tracks = tracks + self.pad_value
 
-        if self.grayscale:
-            transform = transforms.Grayscale()
-            video = transform(video)
-            video = video.repeat(1, 1, 3, 1, 1)
+        # Convert opacity to integer value
+        opacity_value = min(max(int(point_opacity * 255), 0), 255)
 
-        res_video = self.draw_tracks_on_video(
-            video=video,
-            tracks=tracks,
+        # 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,
             visibility=visibility,
-            segm_mask=segm_mask,
-            gt_tracks=gt_tracks,
+            segmentation=segmentation,
             query_frame=query_frame,
-            compensate_for_camera_motion=compensate_for_camera_motion,
-            color_alpha=color_alpha,
+            opacity=opacity_value,
         )
+
+        # Save video if requested
         if save_video:
-            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]
+            self.save_video(output_video, filename=filename)
+
+        return output_video
 
-            # Prepare the video file path
-            save_path = os.path.join(self.save_dir, f"{filename}.mp4")
+    def save_video(self, video: torch.Tensor, filename: str):
+        """Save video tensor as MP4 file.
 
-            # Create a writer object
-            video_writer = imageio.get_writer(save_path, fps=self.fps)
+        Args:
+            video: Video tensor of shape (B,T,C,H,W)
+            filename: Output filename without extension
+        """
+        os.makedirs(self.output_dir, exist_ok=True)
 
-            # Write frames to the video file
-            for frame in wide_list[2:-1]:
-                video_writer.append_data(frame)
+        # Extract frames from video tensor
+        frames = [
+            frame[0].permute(1, 2, 0).cpu().numpy().astype(np.uint8)
+            for frame in video.unbind(1)
+        ]
 
-            video_writer.close()
+        output_path = os.path.join(self.output_dir, f"{filename}.mp4")
 
-            print(f"Video saved to {save_path}")
+        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)
+
+            print(f"Successfully saved video to {output_path}")
+
+        except Exception as e:
+            print(f"Error saving video to {output_path}: {str(e)}")
 
     def draw_tracks_on_video(
         self,
         video: torch.Tensor,
         tracks: torch.Tensor,
         visibility: torch.Tensor = None,
-        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].long().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(),
-                )
-                norm = plt.Normalize(y_min, y_max)
-                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:
-            if self.mode == "rainbow":
-                vector_colors[:, segm_mask <= 0, :] = 255
+        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
+        )
 
-                y_min, y_max = (
-                    tracks[0, segm_mask > 0, 1].min(),
-                    tracks[0, segm_mask > 0, 1].max(),
+        # 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
                 )
-                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
+                color = (
+                    np.array(
+                        self.color_mapper(norm(tracks[query_idx, point_idx, 1]))[:3]
+                    )[None]
+                    * 255
                 )
-                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,
+                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
                 )
-                if gt_tracks is not None:
-                    res_video[t] = self._draw_gt_tracks(
-                        res_video[t], gt_tracks[first_ind : t + 1]
-                    )
+                track_colors[frame_idx] = np.repeat(color, num_points, axis=0)
+
+        return track_colors.astype(np.uint8)
 
-        #  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 = (int(coord[0]) + length, int(coord[1]) + length)
-                            coord_x = (int(coord[0]) - length, int(coord[1]) - length)
-                            rgb = draw_line(
-                                img,
-                                coord_y,
-                                coord_x,
-                                vector_colors[t, i].astype(int),
-                                self.linewidth,
-                            )
-                            coord_y = (int(coord[0]) - length, int(coord[1]) + length)
-                            coord_x = (int(coord[0]) + length, int(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(
+    def _assign_segmentation_colors(
         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 = (int(tracks[s, i, 0]), int(tracks[s, i, 1]))
-                coord_x = (int(tracks[s + 1, i, 0]), int(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.tracks_leave_trace > 0:
-                rgb = Image.fromarray(
-                    np.uint8(
-                        add_weighted(
-                            np.array(rgb), alpha, np.array(original), 1 - alpha, 0
+        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
                         )
-                    )
-                )
-        rgb = np.array(rgb)
-        return rgb
+                        track_colors[:, point_idx] = np.repeat(
+                            color, num_frames, 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
+            )
 
-    def _draw_gt_tracks(
+            # 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
+            )
+
+        return frames
+
+    def _draw_track_segments(
         self,
-        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 = (int(gt_tracks[0]) + length, int(gt_tracks[1]) + length)
-                    coord_x = (int(gt_tracks[0]) - length, int(gt_tracks[1]) - length)
-                    rgb = draw_line(
-                        rgb,
-                        coord_y,
-                        coord_x,
-                        color,
-                        self.linewidth,
+        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],
                     )
-                    coord_y = (int(gt_tracks[0]) - length, int(gt_tracks[1]) + length)
-                    coord_x = (int(gt_tracks[0]) + length, int(gt_tracks[1]) - length)
-                    rgb = draw_line(
-                        rgb,
-                        coord_y,
-                        coord_x,
-                        color,
-                        self.linewidth,
+                    end = (
+                        tracks[segment_idx + 1, point_idx, 0],
+                        tracks[segment_idx + 1, point_idx, 1],
+                    )
+
+                    frame_img = draw_line_segment(
+                        frame_img,
+                        start,
+                        end,
+                        segment_color[point_idx].astype(int),
+                        self.line_width,
                     )
-        rgb = np.array(rgb)
-        return rgb
+
+            if self.track_history_length > 0:
+                frame_img = Image.fromarray(
+                    blend_images(
+                        np.array(frame_img), alpha, np.array(original), 1 - alpha, 0
+                    )
+                )
+
+        return np.array(frame_img)
+
+    def _draw_track_points(
+        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