cotracker 2.0

app.py

@@ -10,6 +10,7 @@ import gradio as gr
 
 from cotracker.utils.visualizer import Visualizer
 
+
 def parse_video(video_file):
     vs = cv2.VideoCapture(video_file)
 
@@ -26,33 +27,33 @@ def parse_video(video_file):
 
 
 def cotracker_demo(
-    input_video, 
-    grid_size: int = 10, 
-    grid_query_frame: int = 0, 
-    backward_tracking: bool = False,
-    tracks_leave_trace: bool = False
-    ):
+    input_video,
+    grid_size: int = 10,
+    grid_query_frame: int = 0,
+    tracks_leave_trace: bool = False,
+):
     load_video = parse_video(input_video)
-    grid_query_frame = min(len(load_video)-1, grid_query_frame)
+    grid_query_frame = min(len(load_video) - 1, grid_query_frame)
     load_video = torch.from_numpy(load_video).permute(0, 3, 1, 2)[None].float()
 
-    model = torch.hub.load("facebookresearch/co-tracker", "cotracker_w8")
+    model = torch.hub.load("facebookresearch/co-tracker", "cotracker2_online")
 
     if torch.cuda.is_available():
         model = model.cuda()
         load_video = load_video.cuda()
-    pred_tracks, pred_visibility = model(
-        load_video, 
-        grid_size=grid_size, 
-        grid_query_frame=grid_query_frame, 
-        backward_tracking=backward_tracking
-        )
+
+    model(video_chunk=load_video, is_first_step=True, grid_size=grid_size)
+    for ind in range(0, load_video.shape[1] - model.step, model.step):
+        pred_tracks, pred_visibility = model(
+            video_chunk=load_video[:, ind : ind + model.step * 2]
+        )  # B T N 2,  B T N 1
+
     linewidth = 2
     if grid_size < 10:
         linewidth = 4
     elif grid_size < 20:
         linewidth = 3
-        
+
     vis = Visualizer(
         save_dir=os.path.join(os.path.dirname(__file__), "results"),
         grayscale=False,
@@ -60,7 +61,7 @@ def cotracker_demo(
         fps=10,
         linewidth=linewidth,
         show_first_frame=5,
-        tracks_leave_trace= -1 if tracks_leave_trace else 0,
+        tracks_leave_trace=-1 if tracks_leave_trace else 0,
     )
     import time
 
@@ -70,53 +71,52 @@ def cotracker_demo(
     filename = str(current_milli_time())
     vis.visualize(
         load_video.cpu(),
-        tracks=pred_tracks.cpu(), 
+        tracks=pred_tracks.cpu(),
         visibility=pred_visibility.cpu(),
-        filename=filename,
+        filename=f"{filename}_pred_track",
         query_frame=grid_query_frame,
-        )
+    )
     return os.path.join(
         os.path.dirname(__file__), "results", f"{filename}_pred_track.mp4"
     )
 
+
 apple = os.path.join(os.path.dirname(__file__), "videos", "apple.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")
 
 app = gr.Interface(
-    title = "🎨 CoTracker: It is Better to Track Together",
-    description = "<div style='text-align: left;'> \
+    title="🎨 CoTracker: It is Better to Track Together",
+    description="<div style='text-align: left;'> \
     <p>Welcome to <a href='http://co-tracker.github.io' target='_blank'>CoTracker</a>! This space demonstrates point (pixel) tracking in videos. \
     Points are sampled on a regular grid and are tracked jointly. </p> \
     <p> To get started, simply upload your <b>.mp4</b> video in landscape orientation or click on one of the example videos to load them. The shorter the video, the faster the processing. We recommend submitting short videos of length <b>2-7 seconds</b>.</p> \
     <ul style='display: inline-block; text-align: left;'> \
         <li>The total number of grid points is the square of <b>Grid Size</b>.</li> \
         <li>To specify the starting frame for tracking, adjust <b>Grid Query Frame</b>. Tracks will be visualized only after the selected frame.</li> \
-        <li>Use <b>Backward Tracking</b> to track points from the selected frame in both directions.</li> \
         <li>Check <b>Visualize Track Traces</b> to visualize traces of all the tracked points. </li> \
     </ul> \
     <p style='text-align: left'>For more details, check out our <a href='https://github.com/facebookresearch/co-tracker' target='_blank'>GitHub Repo</a> ⭐</p> \
     </div>",
-           
     fn=cotracker_demo,
     inputs=[
         gr.Video(type="file", label="Input video", interactive=True),
         gr.Slider(minimum=1, maximum=30, step=1, value=10, label="Grid Size"),
         gr.Slider(minimum=0, maximum=30, step=1, default=0, label="Grid Query Frame"),
-        gr.Checkbox(label="Backward Tracking"),
         gr.Checkbox(label="Visualize Track Traces"),
     ],
     outputs=gr.Video(label="Video with predicted tracks"),
     examples=[
-        [ apple, 10, 0, False, False ],
-        [ apple, 20, 30, True, False ],
-        [ bear, 10, 0, False, False ],
-        [ paragliding, 10, 0, False, False ],
-        [ paragliding_launch, 10, 0, False, False ],
+        [apple, 10, 0, False, False],
+        [apple, 20, 30, True, False],
+        [bear, 10, 0, False, False],
+        [paragliding, 10, 0, False, False],
+        [paragliding_launch, 10, 0, False, False],
     ],
     cache_examples=True,
     allow_flagging=False,
-
 )
 app.queue(max_size=20, concurrency_count=2).launch(debug=True)

requirements.txt

@@ -1,11 +1,7 @@
-einops
-timm
-tqdm
-opencv-python
 matplotlib
-moviepy
 flow_vis
 imutils
 numpy
+imageio
 gradio
 git+https://github.com/facebookresearch/co-tracker.git