Initial commit

README.md

@@ -0,0 +1,24 @@
+---
+title: CoTracker
+emoji: 🎨
+colorFrom: yellow
+colorTo: pink
+sdk: gradio
+sdk_version: 3.41.2
+app_file: app.py
+pinned: false
+license: cc-by-nc-4.0
+---
+
+This is a demo for ["CoTracker: It is Better to Track Together"](https://arxiv.org/abs/2307.07635)
+
+```
+@misc{karaev2023cotracker,
+      title={CoTracker: It is Better to Track Together}, 
+      author={Nikita Karaev and Ignacio Rocco and Benjamin Graham and Natalia Neverova and Andrea Vedaldi and Christian Rupprecht},
+      year={2023},
+      eprint={2307.07635},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```

app.py

@@ -0,0 +1,122 @@
+import os
+import cv2
+import imutils
+import torch
+import timm
+import einops
+import tqdm
+import numpy as np
+import gradio as gr
+
+from cotracker.utils.visualizer import Visualizer
+
+def parse_video(video_file):
+    vs = cv2.VideoCapture(video_file)
+
+    frames = []
+    while True:
+        (gotit, frame) = vs.read()
+        if frame is not None:
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frames.append(frame)
+        if not gotit:
+            break
+
+    return np.stack(frames)
+
+
+def cotracker_demo(
+    input_video, 
+    grid_size: int = 10, 
+    grid_query_frame: int = 0, 
+    backward_tracking: bool = False,
+    tracks_leave_trace: bool = False
+    ):
+    load_video = parse_video(input_video)
+    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")
+
+    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
+        )
+    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,
+        pad_value=100,
+        fps=10,
+        linewidth=linewidth,
+        show_first_frame=5,
+        tracks_leave_trace= -1 if tracks_leave_trace else 0,
+    )
+    import time
+
+    def current_milli_time():
+        return round(time.time() * 1000)
+
+    filename = str(current_milli_time())
+    vis.visualize(
+        load_video.cpu(),
+        tracks=pred_tracks.cpu(), 
+        visibility=pred_visibility.cpu(),
+        filename=filename,
+        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 = 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;'> \
+    <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 ],
+    ],
+    cache_examples=False,
+    allow_flagging=False,
+
+)
+app.queue(max_size=20, concurrency_count=2).launch(debug=True)

pre-requirements.txt

@@ -0,0 +1,2 @@
+torch==1.13.0
+torchvision==0.14.0

requirements.txt

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