From 153b6a3052ee664805eb49990143da72ed8b082a Mon Sep 17 00:00:00 2001 From: Niels Date: Fri, 19 Jul 2024 21:50:30 +0200 Subject: [PATCH 1/2] Add first draft --- cotracker/models/core/cotracker/cotracker.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/cotracker/models/core/cotracker/cotracker.py b/cotracker/models/core/cotracker/cotracker.py index 53178fbe..64416833 100644 --- a/cotracker/models/core/cotracker/cotracker.py +++ b/cotracker/models/core/cotracker/cotracker.py @@ -8,6 +8,8 @@ import torch.nn as nn import torch.nn.functional as F +from huggingface_hub import PyTorchModelHubMixin + from cotracker.models.core.model_utils import sample_features4d, sample_features5d from cotracker.models.core.embeddings import ( get_2d_embedding, @@ -26,7 +28,7 @@ torch.manual_seed(0) -class CoTracker2(nn.Module): +class CoTracker2(nn.Module, PyTorchModelHubMixin): def __init__( self, window_len=8, From 91d2f800c3451d8deea5bb036aa2b5992ba26b94 Mon Sep 17 00:00:00 2001 From: Niels Date: Fri, 19 Jul 2024 22:23:03 +0200 Subject: [PATCH 2/2] Add metadata --- cotracker/models/core/cotracker/cotracker.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/cotracker/models/core/cotracker/cotracker.py b/cotracker/models/core/cotracker/cotracker.py index 64416833..b97d66df 100644 --- a/cotracker/models/core/cotracker/cotracker.py +++ b/cotracker/models/core/cotracker/cotracker.py @@ -28,7 +28,11 @@ torch.manual_seed(0) -class CoTracker2(nn.Module, PyTorchModelHubMixin): +class CoTracker2(nn.Module, PyTorchModelHubMixin, + library_name="co-tracker", + repo_url="https://github.com/facebookresearch/co-tracker", + license="cc-by-nc-4.0", + tags=["object-tracking"]): def __init__( self, window_len=8,