Update visualizer.py

visualizer.py

@@ -255,7 +255,7 @@ class Visualizer:
                     curr_tracks = curr_tracks[:, segm_mask > 0]
                     curr_colors = curr_colors[:, segm_mask > 0]
 
-                res_video[t] = self._draw_pred_tracks(
+                res_video[t] = self._draw_gt_tracks(
                     res_video[t],
                     curr_tracks,
                     curr_colors,
@@ -331,12 +331,18 @@ class Visualizer:
         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
-        color = np.array((211, 0, 0))
+        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: