cotracker/gradio_demo/app.py

102 lines
3.7 KiB
Python
Raw Normal View History

2024-01-05 17:17:50 +01:00
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
2023-09-05 14:04:51 +02:00
import os
import torch
import gradio as gr
from cotracker.utils.visualizer import Visualizer, read_video_from_path
def cotracker_demo(
2023-12-27 13:54:02 +01:00
input_video,
grid_size: int = 10,
grid_query_frame: int = 0,
tracks_leave_trace: bool = False,
):
2023-09-05 14:04:51 +02:00
load_video = read_video_from_path(input_video)
2023-12-27 13:54:02 +01:00
grid_query_frame = min(len(load_video) - 1, grid_query_frame)
2023-09-05 14:04:51 +02:00
load_video = torch.from_numpy(load_video).permute(0, 3, 1, 2)[None].float()
2023-12-27 13:54:02 +01:00
model = torch.hub.load("facebookresearch/co-tracker", "cotracker2_online")
2023-09-05 14:04:51 +02:00
if torch.cuda.is_available():
model = model.cuda()
load_video = load_video.cuda()
2023-12-27 13:54:02 +01:00
2024-01-05 17:17:50 +01:00
model(
video_chunk=load_video,
is_first_step=True,
grid_size=grid_size,
grid_query_frame=grid_query_frame,
)
2023-12-27 13:54:02 +01:00
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
2023-09-05 14:04:51 +02:00
linewidth = 2
if grid_size < 10:
linewidth = 4
elif grid_size < 20:
linewidth = 3
2023-12-27 13:54:02 +01:00
2023-09-05 14:04:51 +02:00
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,
2023-12-27 13:54:02 +01:00
tracks_leave_trace=-1 if tracks_leave_trace else 0,
2023-09-05 14:04:51 +02:00
)
import time
def current_milli_time():
return round(time.time() * 1000)
filename = str(current_milli_time())
vis.visualize(
load_video,
2023-12-27 13:54:02 +01:00
tracks=pred_tracks,
2023-09-05 14:04:51 +02:00
visibility=pred_visibility,
2023-12-27 14:27:56 +01:00
filename=f"{filename}_pred_track",
2023-09-05 14:04:51 +02:00
query_frame=grid_query_frame,
)
2023-12-27 13:54:02 +01:00
return os.path.join(os.path.dirname(__file__), "results", f"{filename}_pred_track.mp4")
2023-09-05 14:04:51 +02:00
app = gr.Interface(
2023-12-27 13:54:02 +01:00
title="🎨 CoTracker: It is Better to Track Together",
description="<div style='text-align: left;'> \
2023-09-05 14:04:51 +02:00
<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>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(label="Input video", interactive=True),
2023-09-05 14:04:51 +02:00
gr.Slider(minimum=1, maximum=30, step=1, value=10, label="Grid Size"),
gr.Slider(minimum=0, maximum=30, step=1, value=0, label="Grid Query Frame"),
2023-09-05 14:04:51 +02:00
gr.Checkbox(label="Visualize Track Traces"),
],
outputs=gr.Video(label="Video with predicted tracks"),
examples=[
2023-12-27 13:54:02 +01:00
["./assets/apple.mp4", 20, 0, False, False],
["./assets/apple.mp4", 10, 30, True, False],
2023-09-05 14:04:51 +02:00
],
2023-12-27 13:54:02 +01:00
cache_examples=False,
2023-09-05 14:04:51 +02:00
)
2023-12-27 14:27:56 +01:00
app.launch(share=True)