2023-07-21 16:30:44 +02: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.
|
|
|
|
|
|
|
|
import torch
|
|
|
|
|
2023-12-27 13:54:02 +01:00
|
|
|
_COTRACKER_URL = "https://huggingface.co/facebook/cotracker/resolve/main/cotracker2.pth"
|
2023-07-21 16:30:44 +02:00
|
|
|
|
|
|
|
|
2023-12-27 13:54:02 +01:00
|
|
|
def _make_cotracker_predictor(*, pretrained: bool = True, online=False, **kwargs):
|
|
|
|
if online:
|
|
|
|
from cotracker.predictor import CoTrackerOnlinePredictor
|
2023-07-21 16:30:44 +02:00
|
|
|
|
2023-12-27 13:54:02 +01:00
|
|
|
predictor = CoTrackerOnlinePredictor(checkpoint=None)
|
|
|
|
else:
|
|
|
|
from cotracker.predictor import CoTrackerPredictor
|
2023-07-21 16:30:44 +02:00
|
|
|
|
2023-12-27 13:54:02 +01:00
|
|
|
predictor = CoTrackerPredictor(checkpoint=None)
|
2023-07-21 16:30:44 +02:00
|
|
|
if pretrained:
|
2023-12-27 13:54:02 +01:00
|
|
|
state_dict = torch.hub.load_state_dict_from_url(_COTRACKER_URL, map_location="cpu")
|
2023-07-21 16:30:44 +02:00
|
|
|
predictor.model.load_state_dict(state_dict)
|
|
|
|
return predictor
|
|
|
|
|
|
|
|
|
2023-12-27 13:54:02 +01:00
|
|
|
def cotracker2(*, pretrained: bool = True, **kwargs):
|
2023-07-21 16:30:44 +02:00
|
|
|
"""
|
2023-12-27 13:54:02 +01:00
|
|
|
CoTracker2 with stride 4 and window length 8. Can track up to 265*265 points jointly.
|
2023-07-21 16:30:44 +02:00
|
|
|
"""
|
2023-12-27 13:54:02 +01:00
|
|
|
return _make_cotracker_predictor(pretrained=pretrained, online=False, **kwargs)
|
|
|
|
|
|
|
|
|
|
|
|
def cotracker2_online(*, pretrained: bool = True, **kwargs):
|
|
|
|
"""
|
|
|
|
Online CoTracker2 with stride 4 and window length 8. Can track up to 265*265 points jointly.
|
|
|
|
"""
|
|
|
|
return _make_cotracker_predictor(pretrained=pretrained, online=True, **kwargs)
|