nikkar commited on
Commit
2aba93c
0 Parent(s):

Initial commit

Browse files
README.md ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: CoTracker
3
+ emoji: 🎨
4
+ colorFrom: yellow
5
+ colorTo: pink
6
+ sdk: gradio
7
+ sdk_version: 3.41.2
8
+ app_file: app.py
9
+ pinned: false
10
+ license: cc-by-nc-4.0
11
+ ---
12
+
13
+ This is a demo for ["CoTracker: It is Better to Track Together"](https://arxiv.org/abs/2307.07635)
14
+
15
+ ```
16
+ @misc{karaev2023cotracker,
17
+ title={CoTracker: It is Better to Track Together},
18
+ author={Nikita Karaev and Ignacio Rocco and Benjamin Graham and Natalia Neverova and Andrea Vedaldi and Christian Rupprecht},
19
+ year={2023},
20
+ eprint={2307.07635},
21
+ archivePrefix={arXiv},
22
+ primaryClass={cs.CV}
23
+ }
24
+ ```
app.py ADDED
@@ -0,0 +1,122 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import cv2
3
+ import imutils
4
+ import torch
5
+ import timm
6
+ import einops
7
+ import tqdm
8
+ import numpy as np
9
+ import gradio as gr
10
+
11
+ from cotracker.utils.visualizer import Visualizer
12
+
13
+ def parse_video(video_file):
14
+ vs = cv2.VideoCapture(video_file)
15
+
16
+ frames = []
17
+ while True:
18
+ (gotit, frame) = vs.read()
19
+ if frame is not None:
20
+ frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
21
+ frames.append(frame)
22
+ if not gotit:
23
+ break
24
+
25
+ return np.stack(frames)
26
+
27
+
28
+ def cotracker_demo(
29
+ input_video,
30
+ grid_size: int = 10,
31
+ grid_query_frame: int = 0,
32
+ backward_tracking: bool = False,
33
+ tracks_leave_trace: bool = False
34
+ ):
35
+ load_video = parse_video(input_video)
36
+ grid_query_frame = min(len(load_video)-1, grid_query_frame)
37
+ load_video = torch.from_numpy(load_video).permute(0, 3, 1, 2)[None].float()
38
+
39
+ model = torch.hub.load("facebookresearch/co-tracker", "cotracker_w8")
40
+
41
+ if torch.cuda.is_available():
42
+ model = model.cuda()
43
+ load_video = load_video.cuda()
44
+ pred_tracks, pred_visibility = model(
45
+ load_video,
46
+ grid_size=grid_size,
47
+ grid_query_frame=grid_query_frame,
48
+ backward_tracking=backward_tracking
49
+ )
50
+ linewidth = 2
51
+ if grid_size < 10:
52
+ linewidth = 4
53
+ elif grid_size < 20:
54
+ linewidth = 3
55
+
56
+ vis = Visualizer(
57
+ save_dir=os.path.join(os.path.dirname(__file__), "results"),
58
+ grayscale=False,
59
+ pad_value=100,
60
+ fps=10,
61
+ linewidth=linewidth,
62
+ show_first_frame=5,
63
+ tracks_leave_trace= -1 if tracks_leave_trace else 0,
64
+ )
65
+ import time
66
+
67
+ def current_milli_time():
68
+ return round(time.time() * 1000)
69
+
70
+ filename = str(current_milli_time())
71
+ vis.visualize(
72
+ load_video.cpu(),
73
+ tracks=pred_tracks.cpu(),
74
+ visibility=pred_visibility.cpu(),
75
+ filename=filename,
76
+ query_frame=grid_query_frame,
77
+ )
78
+ return os.path.join(
79
+ os.path.dirname(__file__), "results", f"{filename}_pred_track.mp4"
80
+ )
81
+
82
+ apple = os.path.join(os.path.dirname(__file__), "videos", "apple.mp4")
83
+ bear = os.path.join(os.path.dirname(__file__), "videos", "bear.mp4")
84
+ paragliding_launch = os.path.join(os.path.dirname(__file__), "videos", "paragliding-launch.mp4")
85
+ paragliding = os.path.join(os.path.dirname(__file__), "videos", "paragliding.mp4")
86
+
87
+ app = gr.Interface(
88
+ title = "🎨 CoTracker: It is Better to Track Together",
89
+ description = "<div style='text-align: left;'> \
90
+ <p>Welcome to <a href='http://co-tracker.github.io' target='_blank'>CoTracker</a>! This space demonstrates point (pixel) tracking in videos. \
91
+ Points are sampled on a regular grid and are tracked jointly. </p> \
92
+ <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> \
93
+ <ul style='display: inline-block; text-align: left;'> \
94
+ <li>The total number of grid points is the square of <b>Grid Size</b>.</li> \
95
+ <li>To specify the starting frame for tracking, adjust <b>Grid Query Frame</b>. Tracks will be visualized only after the selected frame.</li> \
96
+ <li>Use <b>Backward Tracking</b> to track points from the selected frame in both directions.</li> \
97
+ <li>Check <b>Visualize Track Traces</b> to visualize traces of all the tracked points. </li> \
98
+ </ul> \
99
+ <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> \
100
+ </div>",
101
+
102
+ fn=cotracker_demo,
103
+ inputs=[
104
+ gr.Video(type="file", label="Input video", interactive=True),
105
+ gr.Slider(minimum=1, maximum=30, step=1, value=10, label="Grid Size"),
106
+ gr.Slider(minimum=0, maximum=30, step=1, default=0, label="Grid Query Frame"),
107
+ gr.Checkbox(label="Backward Tracking"),
108
+ gr.Checkbox(label="Visualize Track Traces"),
109
+ ],
110
+ outputs=gr.Video(label="Video with predicted tracks"),
111
+ examples=[
112
+ [ apple, 10, 0, False, False ],
113
+ [ apple, 20, 30, True, False ],
114
+ [ bear, 10, 0, False, False ],
115
+ [ paragliding, 10, 0, False, False ],
116
+ [ paragliding_launch, 10, 0, False, False ],
117
+ ],
118
+ cache_examples=False,
119
+ allow_flagging=False,
120
+
121
+ )
122
+ app.queue(max_size=20, concurrency_count=2).launch(debug=True)
pre-requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ torch==1.13.0
2
+ torchvision==0.14.0
requirements.txt ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ einops
2
+ timm
3
+ tqdm
4
+ opencv-python
5
+ matplotlib
6
+ moviepy
7
+ flow_vis
8
+ imutils
9
+ numpy
10
+ git+https://github.com/facebookresearch/co-tracker.git
videos/apple.mp4 ADDED
Binary file (726 kB). View file
 
videos/bear.mp4 ADDED
Binary file (894 kB). View file
 
videos/paragliding-launch.mp4 ADDED
Binary file (829 kB). View file
 
videos/paragliding.mp4 ADDED
Binary file (438 kB). View file