33 lines
978 B
Python
33 lines
978 B
Python
# 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 torch
|
|
|
|
dependencies = ["torch", "einops", "timm", "tqdm"]
|
|
|
|
_COTRACKER_URL = (
|
|
"https://dl.fbaipublicfiles.com/cotracker/cotracker_stride_4_wind_8.pth"
|
|
)
|
|
|
|
|
|
def _make_cotracker_predictor(*, pretrained: bool = True, **kwargs):
|
|
from cotracker.predictor import CoTrackerPredictor
|
|
|
|
predictor = CoTrackerPredictor(checkpoint=None)
|
|
if pretrained:
|
|
state_dict = torch.hub.load_state_dict_from_url(
|
|
_COTRACKER_URL, map_location="cpu"
|
|
)
|
|
predictor.model.load_state_dict(state_dict)
|
|
return predictor
|
|
|
|
|
|
def cotracker_w8(*, pretrained: bool = True, **kwargs):
|
|
"""
|
|
CoTracker model with stride 4 and window length 8. (The main model from the paper)
|
|
"""
|
|
return _make_cotracker_predictor(pretrained=pretrained, **kwargs)
|