File size: 3,941 Bytes
2aba93c
 
 
 
 
 
 
 
 
d6d3990
2aba93c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d6d3990
 
 
 
2aba93c
 
 
d6d3990
2aba93c
 
 
 
d6d3990
6393e70
d6d3990
 
 
 
 
2aba93c
 
 
 
 
d6d3990
2aba93c
 
 
 
 
 
 
d6d3990
2aba93c
 
 
 
 
 
 
 
 
d6d3990
2aba93c
d6d3990
 
2aba93c
 
 
 
d6d3990
2aba93c
 
d6d3990
 
 
2aba93c
 
 
d6d3990
 
2aba93c
 
 
 
 
 
 
 
 
 
 
 
d515d68
2aba93c
 
 
 
6393e70
 
 
 
 
2aba93c
737c008
2aba93c
 
737c008
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
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]
        )  # 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,
        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)