|
import os |
|
import cv2 |
|
import imutils |
|
import torch |
|
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, |
|
tracks_leave_trace: bool = False, |
|
): |
|
load_video = parse_video(input_video) |
|
load_video = torch.from_numpy(load_video).permute(0, 3, 1, 2)[None].float() |
|
|
|
model = torch.hub.load("facebookresearch/co-tracker", "cotracker2_online") |
|
|
|
if torch.cuda.is_available(): |
|
model = model.cuda() |
|
load_video = load_video.cuda() |
|
|
|
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] |
|
) |
|
|
|
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=f"{filename}_pred_track", |
|
) |
|
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>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=10, maximum=80, step=1, value=10, label="Grid Size"), |
|
gr.Checkbox(label="Visualize Track Traces"), |
|
], |
|
outputs=gr.Video(label="Video with predicted tracks"), |
|
examples=[ |
|
[apple, 30, False], |
|
[apple, 10, True], |
|
[bear, 10, False], |
|
[paragliding, 10, False], |
|
[paragliding_launch, 10, False], |
|
], |
|
cache_examples=True, |
|
allow_flagging=False, |
|
) |
|
app.queue(max_size=20, concurrency_count=2).launch(debug=True) |
|
|
