added rate slider

app.py

@@ -103,8 +103,8 @@ def paint_point_track_gpu_scatter(
         point_tracks: np.ndarray,
         visibles: np.ndarray,
         colormap: Optional[List[Tuple[int, int, int]]] = None,
-        radius: int = 1,
-        sharpness: float = 0.15,
+        rate: int = 1,
+        # sharpness: float = 0.1,
 ) -> np.ndarray:
     print('starting vis')
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -117,6 +117,20 @@ def paint_point_track_gpu_scatter(
     if colormap is None:
         colormap = get_colors(P)
     colors = torch.tensor(colormap, dtype=torch.float32, device=device)  # [P,3]
+
+    if rate==1:
+        radius = 1
+    elif rate==2:
+        radius = 1
+    elif rate== 4:
+        radius = 2
+    elif rate== 8:
+        radius = 4
+    else:
+        radius = 6
+    # radius = max(1, int(np.sqrt(rate)))
+    sharpness = 0.15 + 0.05 * np.log2(rate)
+    
     D = radius * 2 + 1
     y = torch.arange(D, device=device).float()[:, None] - radius
     x = torch.arange(D, device=device).float()[None, :] - radius
@@ -160,7 +174,7 @@ def paint_point_track_gpu_scatter(
         # frames_t[t] = frames_t[t] * (1 - weight) + accum
 
         # alpha = weight.clamp(0, 1)
-        # alpha = weight.clamp(0, 1) * 0.75  # transparency
+        # alpha = weight.clamp(0, 1) * 0.9  # transparency
         alpha = weight.clamp(0, 1)  # transparency
         accum = accum / (weight + 1e-6)  # [3, H, W]
         frames_t[t] = frames_t[t] * (1 - alpha) + accum * alpha
@@ -491,6 +505,25 @@ FRAME_LIMIT = 600 # Limit the number of frames to process
 def choose_frame(frame_num, video_preview_array):
     return video_preview_array[int(frame_num)]
 
+def choose_rate1(video_preview, video_fps, tracks, visibs):
+    return choose_rate(1, video_preview, video_fps, tracks, visibs)
+def choose_rate2(video_preview, video_fps, tracks, visibs):
+    return choose_rate(2, video_preview, video_fps, tracks, visibs)
+def choose_rate4(video_preview, video_fps, tracks, visibs):
+    return choose_rate(4, video_preview, video_fps, tracks, visibs)
+def choose_rate8(video_preview, video_fps, tracks, visibs):
+    return choose_rate(8, 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):
+    print('rate', rate)
+    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 video_preview_array[int(frame_num)]
 
 def preprocess_video_input(video_path):
     video_arr = mediapy.read_video(video_path)
@@ -509,7 +542,7 @@ def preprocess_video_input(video_path):
     if height*width > 768*1024:
         new_height = new_height*3//4
         new_width = new_width*3//4
-    new_height, new_width = new_height//8 * 8, new_width//8 * 8 # make it divisible by 8, partly to satisfy ffmpeg
+    new_height, new_width = new_height//16 * 16, new_width//16 * 16 # make it divisible by 16, partly to satisfy ffmpeg
         
         
     preview_video = mediapy.resize_video(video_arr, (new_height, new_width))
@@ -541,8 +574,56 @@ def preprocess_video_input(video_path):
         gr.update(interactive=interactive),
         gr.update(interactive=interactive),
         gr.update(interactive=True),
+        # gr.update(interactive=True),
+        # gr.update(interactive=True),
+        # gr.update(interactive=True),
+        # gr.update(interactive=True),
     )
 
+
+def paint_video(video_preview, video_fps, tracks, visibs, rate=1):
+    print('video_preview', video_preview.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)
+    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"
+    video_path = os.path.join(os.path.dirname(__file__), "tmp")
+    video_file_path = os.path.join(video_path, video_file_name)
+    os.makedirs(video_path, exist_ok=True)
+    if False:
+        mediapy.write_video(video_file_path, painted_video, fps=video_fps)
+    else:
+        for ti in range(T):
+            temp_out_f = '%s/%03d.jpg' % (video_path, ti)
+            # temp_out_f = '%s/%03d.png' % (video_path, ti)
+            im = PIL.Image.fromarray(painted_video[ti])
+            # im.save(temp_out_f, "PNG", subsampling=0, quality=80)
+            im.save(temp_out_f)
+            print('saved', temp_out_f)
+        # os.system('/usr/bin/ffmpeg -y -hide_banner -loglevel error -f image2 -framerate %d -pattern_type glob -i "%s/*.png" -c:v libx264 -crf 20 -pix_fmt yuv420p %s' % (video_fps, video_path, video_file_path))
+        os.system('/usr/bin/ffmpeg -y -hide_banner -loglevel error -f image2 -framerate %d -pattern_type glob -i "%s/*.jpg" -c:v libx264 -crf 20 -pix_fmt yuv420p %s' % (video_fps, video_path, video_file_path))
+        print('saved', video_file_path)
+        for ti in range(T):
+            # temp_out_f = '%s/%03d.png' % (video_path, ti)
+            temp_out_f = '%s/%03d.jpg' % (video_path, ti)
+            os.remove(temp_out_f)
+            print('deleted', temp_out_f)
+    return video_file_path
+
+
 @spaces.GPU
 def track(
     video_preview,
@@ -696,78 +777,86 @@ def track(
     # traj_maps_e = traj_maps_e[:,:,:,::4,::4] # subsample
     # visconf_maps_e = visconf_maps_e[:,:,:,::4,::4] # subsample
     
-    traj_maps_e = traj_maps_e[:,:,:,::2,::2] # subsample
-    visconf_maps_e = visconf_maps_e[:,:,:,::2,::2] # subsample
+    # traj_maps_e = traj_maps_e[:,:,:,::2,::2] # subsample
+    # visconf_maps_e = visconf_maps_e[:,:,:,::2,::2] # subsample
 
     tracks = traj_maps_e.permute(0,3,4,1,2).reshape(-1,T,2).numpy()
     visibs = visconf_maps_e.permute(0,3,4,1,2).reshape(-1,T,2)[:,:,0].numpy()
     confs = visconf_maps_e.permute(0,3,4,1,2).reshape(-1,T,2)[:,:,0].numpy()
-
-    visibs = (visibs * confs) > 0.2 # N,T
+    visibs = (visibs * confs) > 0.3 # N,T
     # visibs = (confs) > 0.1 # N,T
     
 
     # sc = (np.array([video_preview.shape[2], video_preview.shape[1]]) / np.array([VIDEO_INPUT_RESO[1], VIDEO_INPUT_RESO[0]])).reshape(1,1,2)
     # print('sc', sc)
     # tracks = tracks * sc
-    
-    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)
 
-    visibs_ = visibs * 1.0
-    visibs_ = visibs_[:,1:] * visibs_[:,:-1]
-    inds = np.sum(visibs_, axis=1) >= min(T//4,8)
-    tracks = tracks[inds]
-    visibs = visibs[inds]
-    colors = colors[inds]
+    return paint_video(video_preview, video_fps, tracks, visibs), tracks, visibs, gr.update(interactive=True, value=1)
+            # gr.update(interactive=True),
+            # gr.update(interactive=True),
+            # gr.update(interactive=True),
+            # gr.update(interactive=True),
+            # gr.update(interactive=True))
+    # # query_count = tracks.shape[0]
     
-    # painted_video = paint_point_track_parallel(video_preview,tracks,visibs,colors)
-    # painted_video = paint_point_track_gpu(video_preview,tracks,visibs,colors)
-    painted_video = paint_point_track_gpu_scatter(video_preview,tracks,visibs,colors)
-    print("7 torch.cuda.memory_allocated: %.1fGB"%(torch.cuda.memory_allocated(0)/1024/1024/1024))
-
-    # save video
-    video_file_name = uuid.uuid4().hex + ".mp4"
-    video_path = os.path.join(os.path.dirname(__file__), "tmp")
-    video_file_path = os.path.join(video_path, video_file_name)
-
-    os.makedirs(video_path, exist_ok=True)
-    if False:
-        mediapy.write_video(video_file_path, painted_video, fps=video_fps)
-    else:
-        for ti in range(T):
-            temp_out_f = '%s/%03d.jpg' % (video_path, ti)
-            # temp_out_f = '%s/%03d.png' % (video_path, ti)
-            im = PIL.Image.fromarray(painted_video[ti])
-            # im.save(temp_out_f, "PNG", subsampling=0, quality=80)
-            im.save(temp_out_f)
-            print('saved', temp_out_f)
-        # os.system('/usr/bin/ffmpeg -y -hide_banner -loglevel error -f image2 -framerate %d -pattern_type glob -i "%s/*.png" -c:v libx264 -crf 20 -pix_fmt yuv420p %s' % (video_fps, video_path, video_file_path))
-        os.system('/usr/bin/ffmpeg -y -hide_banner -loglevel error -f image2 -framerate %d -pattern_type glob -i "%s/*.jpg" -c:v libx264 -crf 20 -pix_fmt yuv420p %s' % (video_fps, video_path, video_file_path))
-        print('saved', video_file_path)
-        for ti in range(T):
-            # temp_out_f = '%s/%03d.png' % (video_path, ti)
-            temp_out_f = '%s/%03d.jpg' % (video_path, ti)
-            os.remove(temp_out_f)
-            print('deleted', temp_out_f)
+    
+    # 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)
 
-    # out_file = tempfile.NamedTemporaryFile(suffix="out.mp4", delete=False)
-    # subprocess.run(f"ffmpeg -y -loglevel quiet -stats -i {painted_video} -c:v libx264 {out_file.name}".split())
+    # # visibs_ = visibs * 1.0
+    # # visibs_ = visibs_[:,1:] * visibs_[:,:-1]
+    # # inds = np.sum(visibs_, axis=1) >= min(T//4,8)
+    # # tracks = tracks[inds]
+    # # visibs = visibs[inds]
+    # # colors = colors[inds]
+    
+    # # painted_video = paint_point_track_parallel(video_preview,tracks,visibs,colors)
+    # # painted_video = paint_point_track_gpu(video_preview,tracks,visibs,colors)
+    # painted_video = paint_point_track_gpu_scatter(video_preview,tracks,visibs,colors)
+    # print("7 torch.cuda.memory_allocated: %.1fGB"%(torch.cuda.memory_allocated(0)/1024/1024/1024))
+
+    # # save video
+    # video_file_name = uuid.uuid4().hex + ".mp4"
+    # video_path = os.path.join(os.path.dirname(__file__), "tmp")
+    # video_file_path = os.path.join(video_path, video_file_name)
+
+    # os.makedirs(video_path, exist_ok=True)
+    # if False:
+    #     mediapy.write_video(video_file_path, painted_video, fps=video_fps)
+    # else:
+    #     for ti in range(T):
+    #         temp_out_f = '%s/%03d.jpg' % (video_path, ti)
+    #         # temp_out_f = '%s/%03d.png' % (video_path, ti)
+    #         im = PIL.Image.fromarray(painted_video[ti])
+    #         # im.save(temp_out_f, "PNG", subsampling=0, quality=80)
+    #         im.save(temp_out_f)
+    #         print('saved', temp_out_f)
+    #     # os.system('/usr/bin/ffmpeg -y -hide_banner -loglevel error -f image2 -framerate %d -pattern_type glob -i "%s/*.png" -c:v libx264 -crf 20 -pix_fmt yuv420p %s' % (video_fps, video_path, video_file_path))
+    #     os.system('/usr/bin/ffmpeg -y -hide_banner -loglevel error -f image2 -framerate %d -pattern_type glob -i "%s/*.jpg" -c:v libx264 -crf 20 -pix_fmt yuv420p %s' % (video_fps, video_path, video_file_path))
+    #     print('saved', video_file_path)
+    #     for ti in range(T):
+    #         # temp_out_f = '%s/%03d.png' % (video_path, ti)
+    #         temp_out_f = '%s/%03d.jpg' % (video_path, ti)
+    #         os.remove(temp_out_f)
+    #         print('deleted', temp_out_f)
+
+    # # out_file = tempfile.NamedTemporaryFile(suffix="out.mp4", delete=False)
+    # # subprocess.run(f"ffmpeg -y -loglevel quiet -stats -i {painted_video} -c:v libx264 {out_file.name}".split())
     
     
 
-    return video_file_path
+    # return video_file_path
 
 
 with gr.Blocks() as demo:
@@ -782,13 +871,17 @@ with gr.Blocks() as demo:
     is_tracked_query = gr.State([])
     query_count = gr.State(0)
 
+    # rate = gr.State([])
+    tracks = gr.State([])
+    visibs = gr.State([])
+
     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 space demonstrates all-pixel tracking in videos.</p> \
     <p>To get started, simply upload your <b>.mp4</b> video, or click on 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 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>!</p> \
-    <p>Initial code for this Gradio app came from LocoTrack and CoTracker.</p> \
+    <p>Initial code for this Gradio app came from LocoTrack and CoTracker -- big thanks to those authors!</p> \
     </div>"
     )
     
@@ -797,25 +890,26 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             with gr.Row():
-                video_in = gr.Video(label="Video Input", format="mp4")
+                video_in = gr.Video(label="Video input", format="mp4")
             with gr.Row():
                 submit = gr.Button("Submit")
         with gr.Column():
             # with gr.Accordion("Sample videos", open=True) as video_in_drawer:
             with gr.Row():
-                dog = os.path.join(os.path.dirname(__file__), "videos", "dog.mp4")
+                butterfly = os.path.join(os.path.dirname(__file__), "videos", "butterfly_800.mp4")
                 monkey = os.path.join(os.path.dirname(__file__), "videos", "monkey_28.mp4")
+                groundbox = os.path.join(os.path.dirname(__file__), "videos", "ground-box-comp.mp4")
                 apple = os.path.join(os.path.dirname(__file__), "videos", "apple.mp4")
+                grasp_sponge_800 = os.path.join(os.path.dirname(__file__), "videos", "grasp_sponge_800.mp4")
+                # dog = os.path.join(os.path.dirname(__file__), "videos", "dog.mp4")
                 bear = os.path.join(os.path.dirname(__file__), "videos", "bear.mp4")
-                paragliding_launch = os.path.join(
-                    os.path.dirname(__file__), "videos", "paragliding-launch.mp4"
-                )
+                paragliding_launch = os.path.join(os.path.dirname(__file__), "videos", "paragliding-launch.mp4")
                 paragliding = os.path.join(os.path.dirname(__file__), "videos", "paragliding.mp4")
                 cat = os.path.join(os.path.dirname(__file__), "videos", "cat.mp4")
                 pillow = os.path.join(os.path.dirname(__file__), "videos", "pillow.mp4")
                 teddy = os.path.join(os.path.dirname(__file__), "videos", "teddy.mp4")
                 backpack = os.path.join(os.path.dirname(__file__), "videos", "backpack.mp4")
-                gr.Examples(examples=[dog, monkey, bear, apple, paragliding, paragliding_launch, cat, pillow, teddy, backpack], 
+                gr.Examples(examples=[butterfly, groundbox, monkey, grasp_sponge_800, bear, apple, paragliding, paragliding_launch, cat, pillow, teddy, backpack], 
                             inputs = [
                                 video_in
                             ],
@@ -828,7 +922,7 @@ with gr.Blocks() as demo:
         with gr.Column():
             with gr.Row():
                 query_frame_slider = gr.Slider(
-                    minimum=0, maximum=100, value=0, step=1, label="Choose Frame", interactive=False)
+                    minimum=0, maximum=100, value=0, step=1, label="Choose frame", interactive=False)
             # with gr.Row():
             #     undo = gr.Button("Undo", interactive=False)
             #     clear_frame = gr.Button("Clear Frame", interactive=False)
@@ -846,16 +940,25 @@ with gr.Blocks() as demo:
                 track_button = gr.Button("Track", interactive=False)
 
         with gr.Column():
+            
             # with gr.Row():
-            #     rate_slider = gr.Slider(
-            #         minimum=1, maximum=16, value=1, step=1, label="Choose subsampling rate", interactive=False)
-            # with gr.Row():
-            output_video = gr.Video(
-                label="Output Video",
-                interactive=False,
-                autoplay=True,
-                loop=True,
-            )
+            #     rate1_button = gr.Button("Subsampling", interactive=False)
+            #     rate2_button = gr.Button("Stride 2", interactive=False)
+            #     rate4_button = gr.Button("Rate 4", interactive=False)
+            #     rate8_button = gr.Button("Rate 8", interactive=False)
+            #     # rate16_button = gr.Button("Rate 16", interactive=False)
+            with gr.Row():
+                # 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():
+                output_video = gr.Video(
+                    label="Output video",
+                    interactive=False,
+                    autoplay=True,
+                    loop=True,
+                )
 
     
 
@@ -892,6 +995,8 @@ with gr.Blocks() as demo:
         queue = False
     )
 
+    
+
     # current_frame.select(
     #     fn = get_point, 
     #     inputs = [
@@ -982,10 +1087,68 @@ with gr.Blocks() as demo:
         ],
         outputs = [
             output_video,
+            tracks,
+            visibs,
+            rate_radio,
+            # rate1_button,
+            # rate2_button,
+            # rate4_button,
+            # rate8_button,
+            # rate16_button,
         ],
         queue = True,
     )
 
+    # rate_slider.change(
+    #     fn = choose_rate,
+    #     inputs = [rate_slider, video_preview, video_fps, tracks, visibs],
+    #     outputs = [
+    #         output_video,
+    #     ],
+    #     queue = False
+    # )
+    rate_radio.change(
+        fn = choose_rate,
+        inputs = [rate_radio, video_preview, video_fps, tracks, visibs],
+        outputs = [
+            output_video,
+        ],
+        queue = False
+    )
+
+    # rate1_button.click(
+    #     fn = choose_rate1,
+    #     inputs = [video_preview, video_fps, tracks, visibs],
+    #     outputs = [output_video],
+    #     queue = False,
+    # )
+    # rate2_button.click(
+    #     fn = choose_rate2,
+    #     inputs = [video_preview, video_fps, tracks, visibs],
+    #     outputs = [output_video],
+    #     queue = False,
+    # )
+    # rate4_button.click(
+    #     fn = choose_rate4,
+    #     inputs = [video_preview, video_fps, tracks, visibs],
+    #     outputs = [output_video],
+    #     queue = False,
+    # )
+    # rate8_button.click(
+    #     fn = choose_rate8,
+    #     inputs = [video_preview, video_fps, tracks, visibs],
+    #     outputs = [output_video],
+    #     queue = False,
+    # )
+    # rate16_button.click(
+    #     fn = choose_rate16,
+    #     inputs = [video_preview, video_fps, tracks, visibs],
+    #     outputs = [output_video],
+    #     queue = False,
+    # )
+    
+    
+