Spaces:
Running
on
L40S
Running
on
L40S
File size: 3,442 Bytes
c705408 |
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 |
# 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.
import glob
import os
import torch
import argparse
import numpy as np
from PIL import Image
from cotracker.utils.visualizer import Visualizer, read_video_from_path
from cotracker.predictor import CoTrackerPredictor
# Unfortunately MPS acceleration does not support all the features we require,
# but we may be able to enable it in the future
DEFAULT_DEVICE = (
# "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
"cuda"
if torch.cuda.is_available()
else "cpu"
)
# if DEFAULT_DEVICE == "mps":
# os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--video_path",
default="./assets/apple.mp4",
help="path to a video",
)
parser.add_argument(
"--mask_path",
default="./assets/apple_mask.png",
help="path to a segmentation mask",
)
parser.add_argument(
"--checkpoint",
default="./checkpoints/cotracker2.pth",
# default=None,
help="CoTracker model parameters",
)
parser.add_argument("--grid_size", type=int, default=10, help="Regular grid size")
parser.add_argument(
"--grid_query_frame",
type=int,
default=0,
help="Compute dense and grid tracks starting from this frame",
)
parser.add_argument(
"--backward_tracking",
action="store_true",
help="Compute tracks in both directions, not only forward",
)
args = parser.parse_args()
if args.checkpoint is not None:
model = CoTrackerPredictor(checkpoint=args.checkpoint)
else:
model = torch.hub.load("facebookresearch/co-tracker", "cotracker2")
model = model.to(DEFAULT_DEVICE)
video_path_list = glob.glob("assets/*.mp4")
# video_path_list = glob.glob("data/vid/*.mp4")
# sort
# video_path_list.sort()
for video_path in video_path_list:
args.video_path = video_path
# load the input video frame by frame
video = read_video_from_path(args.video_path)
# (t, h, w, c) -> (t, c, h, w)
video = torch.from_numpy(video).permute(0, 3, 1, 2)[None].float()
# segm_mask = np.array(Image.open(os.path.join(args.mask_path)))
# segm_mask = torch.from_numpy(segm_mask)[None, None]
video = video.to(DEFAULT_DEVICE)
video = video[:, :200]
with torch.no_grad():
pred_tracks, pred_visibility = model(
video,
grid_size=args.grid_size, # 10
grid_query_frame=args.grid_query_frame, # 0
backward_tracking=args.backward_tracking, # False
# segm_mask=segm_mask,
)
print("computed")
# save a video with predicted tracks
seq_name = os.path.splitext(args.video_path.split("/")[-1])[0]
vis = Visualizer(save_dir="./saved_videos", pad_value=120, linewidth=3)
vis.visualize(
video,
pred_tracks, # (b, f, num_points, 2)
pred_visibility, # (b, f, num_points)
query_frame=0 if args.backward_tracking else args.grid_query_frame,
filename=seq_name,
)
|