cotracker/hubconf.py

39 lines
1.3 KiB
Python
Raw Normal View History

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)