added colormap options

app.py

@@ -517,15 +517,16 @@ def choose_rate8(video_preview, video_fps, tracks, visibs):
 # def choose_rate16(video_preview, video_fps, tracks, visibs):
 #     return choose_rate(16, video_preview, video_fps, tracks, visibs)
 
-def choose_rate(rate, video_preview, video_fps, tracks, visibs):
+def update_vis(rate, cmap, video_preview, query_frame, video_fps, tracks, visibs):
     print('rate', rate)
+    print('cmap', cmap)
     print('video_preview', video_preview.shape)
     T, H, W,_ = video_preview.shape
     tracks_ = tracks.reshape(H,W,T,2)[::rate,::rate].reshape(-1,T,2)
     visibs_ = visibs.reshape(H,W,T)[::rate,::rate].reshape(-1,T)
-    return paint_video(video_preview, video_fps, tracks_, visibs_, rate=rate)
+    return paint_video(video_preview, query_frame, video_fps, tracks_, visibs_, rate=rate, cmap=cmap)
     # return video_preview_array[int(frame_num)]
-
+    
 def preprocess_video_input(video_path):
     video_arr = mediapy.read_video(video_path)
     video_fps = video_arr.metadata.fps
@@ -553,27 +554,15 @@ def preprocess_video_input(video_path):
     preview_video = np.array(preview_video)
     input_video = np.array(input_video)
     
-    interactive = True
-
     return (
         video_arr, # Original video
         preview_video, # Original preview video, resized for faster processing
         preview_video.copy(), # Copy of preview video for visualization
         input_video, # Resized video input for model
-        # None, # video_feature, # Extracted feature
         video_fps, # Set the video FPS
-        # gr.update(open=True), # open/close the video input drawer
-        # tracking_mode, # Set the tracking mode
         preview_video[0], # Set the preview frame to the first frame
-        gr.update(minimum=0, maximum=num_frames - 1, value=0, interactive=interactive), # Set slider interactive
-        [[] for _ in range(num_frames)], # Set query_points to empty
-        [[] for _ in range(num_frames)], # Set query_points_color to empty
-        [[] for _ in range(num_frames)], 
-        0, # Set query count to 0
-        gr.update(interactive=interactive), # Make the buttons interactive
-        gr.update(interactive=interactive),
-        gr.update(interactive=interactive),
-        gr.update(interactive=True),
+        gr.update(minimum=0, maximum=num_frames - 1, value=0, interactive=True), # Set slider interactive
+        gr.update(interactive=True), # make track button interactive
         # gr.update(interactive=True),
         # gr.update(interactive=True),
         # gr.update(interactive=True),
@@ -581,22 +570,30 @@ def preprocess_video_input(video_path):
     )
 
 
-def paint_video(video_preview, video_fps, tracks, visibs, rate=1):
+def paint_video(video_preview, query_frame, video_fps, tracks, visibs, rate=1, cmap="gist_rainbow"):
     print('video_preview', video_preview.shape)
+    print('tracks', tracks.shape)
     T, H, W, _ = video_preview.shape
     query_count = tracks.shape[0]
-    cmap = matplotlib.colormaps.get_cmap("gist_rainbow")
-    query_points_color = [[]]
-    for i in range(query_count):
-        # Choose the color for the point from matplotlib colormap
-        color = cmap(i / float(query_count))
-        color = (int(color[0] * 255), int(color[1] * 255), int(color[2] * 255))
-        query_points_color[0].append(color)
-    # make color array
-    colors = []
-    for frame_colors in query_points_color:
-        colors.extend(frame_colors)
-    colors = np.array(colors)
+    print('cmap', cmap)
+
+    if cmap=="bremm":
+        xy0 = tracks[:,query_frame] # N,2
+        colors = utils.improc.get_2d_colors(xy0, H, W)
+    else:
+        cmap_ = matplotlib.colormaps.get_cmap(cmap)
+        query_points_color = [[]]
+        for i in range(query_count):
+            # Choose the color for the point from matplotlib colormap
+            color = cmap_(i / float(query_count))
+            color = (int(color[0] * 255), int(color[1] * 255), int(color[2] * 255))
+            query_points_color[0].append(color)
+        # make color array
+        colors = []
+        for frame_colors in query_points_color:
+            colors.extend(frame_colors)
+        colors = np.array(colors)
+        
     painted_video = paint_point_track_gpu_scatter(video_preview,tracks,visibs,colors,rate=rate)#=max(rate//2,1))
     # save video
     video_file_name = uuid.uuid4().hex + ".mp4"
@@ -630,9 +627,6 @@ def track(
     video_input, 
     video_fps,
     query_frame,
-    query_points, 
-    query_points_color, 
-    query_count, 
 ):
     # tracking_mode = 'selected'
     # if query_count == 0: 
@@ -788,7 +782,7 @@ def track(
     # print('sc', sc)
     # tracks = tracks * sc
 
-    return paint_video(video_preview, video_fps, tracks, visibs), tracks, visibs, gr.update(interactive=True, value=1)
+    return paint_video(video_preview, query_frame, video_fps, tracks, visibs), tracks, visibs, gr.update(interactive=True), gr.update(interactive=True)
             # gr.update(interactive=True),
             # gr.update(interactive=True),
             # gr.update(interactive=True),
@@ -863,11 +857,6 @@ with gr.Blocks() as demo:
     video_input = gr.State()
     video_fps = gr.State(24)
 
-    query_points = gr.State([])
-    query_points_color = gr.State([])
-    is_tracked_query = gr.State([])
-    query_count = gr.State(0)
-
     # rate = gr.State([])
     tracks = gr.State([])
     visibs = gr.State([])
@@ -875,14 +864,13 @@ with gr.Blocks() as demo:
     gr.Markdown("# ⚡ AllTracker: Efficient Dense Point Tracking at High Resolution")
     gr.Markdown("<div style='text-align: left;'> \
     <p>Welcome to <a href='https://alltracker.github.io/' target='_blank'>AllTracker</a>! This demo runs our model to perform all-pixel tracking in a video of your choice.</p> \
-    <p>To get started, simply upload your <b>.mp4</b> video, or select one of the example videos. The shorter the video, the faster the processing. We recommend submitting videos under 20 seconds long.</p> \
+    <p>To get started, simply upload an mp4, or select one of the example videos. The shorter the video, the faster the processing. We recommend submitting videos under 20 seconds long.</p> \
     <p>After picking a video, click \"Submit\" to load the frames into the app, and optionally choose a query frame (using the slider), and then click \"Track\".</p> \
     <p>For full info on how this works, check out our <a href='https://github.com/aharley/alltracker/' target='_blank'>GitHub repo</a>, or <a href='https://arxiv.org/abs/2506.07310' target='_blank'>paper</a>.</p> \
     <p>Initial code for this Gradio app came from LocoTrack and CoTracker -- big thanks to those authors!</p> \
     </div>"
     )
     
-
     gr.Markdown("## Step 1: Select a video, and click \"Submit\".")
     with gr.Row():
         with gr.Column():
@@ -891,7 +879,6 @@ with gr.Blocks() as demo:
             with gr.Row():
                 submit = gr.Button("Submit")
         with gr.Column():
-            # with gr.Accordion("Sample videos", open=True) as video_in_drawer:
             with gr.Row():
                 butterfly = os.path.join(os.path.dirname(__file__), "videos", "butterfly_800.mp4")
                 monkey = os.path.join(os.path.dirname(__file__), "videos", "monkey_800.mp4")
@@ -951,6 +938,9 @@ with gr.Blocks() as demo:
                 # rate_slider = gr.Slider(
                 #     minimum=1, maximum=16, value=1, step=1, label="Choose subsampling rate", interactive=False)
                 rate_radio = gr.Radio([1, 2, 4, 8, 16], value=1, label="Choose visualization subsampling", interactive=False)
+
+            with gr.Row():
+                cmap_radio = gr.Radio(["gist_rainbow", "rainbow", "jet", "turbo", "bremm"], value="gist_rainbow", label="Choose colormap", interactive=False)
                         
             with gr.Row():
                 output_video = gr.Video(
@@ -971,13 +961,8 @@ with gr.Blocks() as demo:
             video_queried_preview,
             video_input,
             video_fps,
-            # video_in_drawer,
             current_frame,
             query_frame_slider,
-            query_points,
-            query_points_color,
-            is_tracked_query,
-            query_count,
             # undo,
             # clear_frame,
             # clear_all,
@@ -1081,15 +1066,13 @@ with gr.Blocks() as demo:
             video_input,
             video_fps,
             query_frame_slider,
-            query_points,
-            query_points_color,
-            query_count,
         ],
         outputs = [
             output_video,
             tracks,
             visibs,
             rate_radio,
+            cmap_radio,
             # rate1_button,
             # rate2_button,
             # rate4_button,
@@ -1108,8 +1091,16 @@ with gr.Blocks() as demo:
     #     queue = False
     # )
     rate_radio.change(
-        fn = choose_rate,
-        inputs = [rate_radio, video_preview, video_fps, tracks, visibs],
+        fn = update_vis,
+        inputs = [rate_radio, cmap_radio, video_preview, query_frame_slider, video_fps, tracks, visibs],
+        outputs = [
+            output_video,
+        ],
+        queue = False
+    )
+    cmap_radio.change(
+        fn = update_vis,
+        inputs = [rate_radio, cmap_radio, video_preview, query_frame_slider, video_fps, tracks, visibs],
         outputs = [
             output_video,
         ],
@@ -1153,5 +1144,5 @@ with gr.Blocks() as demo:
 
     
 # demo.launch(show_api=False, show_error=True, debug=False, share=False)
-# demo.launch(show_api=False, show_error=True, debug=False, share=True)
-demo.launch(show_api=False, show_error=True, debug=False, share=False)
+demo.launch(show_api=False, show_error=True, debug=False, share=True)
+# demo.launch(show_api=False, show_error=True, debug=False, share=False)

utils/improc.py

@@ -58,7 +58,7 @@ def flow2color(flow, clip=0.0):
     flow = (flow*255.0).type(torch.ByteTensor)
     return flow
 
-COLORMAP_FILE = "./utils/bremm.png"
+COLORMAP_FILE = "./bremm.png"
 class ColorMap2d:
     def __init__(self, filename=None):
         self._colormap_file = filename or COLORMAP_FILE