Minor changes

Author: amanteur

oml/models/__init__.py

@@ -1,3 +1,4 @@
+from oml.models.audio.ecapa_tdnn.extractor import ECAPATDNNExtractor
 from oml.models.meta.projection import ExtractorWithMLP
 from oml.models.meta.siamese import (
     ConcatSiamese,

oml/models/audio/ecapa_tdnn/extractor.py

@@ -82,7 +82,12 @@ def feat_dim(self) -> int:
         return self.model.fc6.weight.shape[0]
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        assert x.ndim == 2, "The model expects input audio to have shape (batch_size, n_samples)"
+        assert x.ndim == 2 or (
+            x.ndim == 3 and x.shape[1] == 1
+        ), "The model expects input audio to have shape (batch_size, n_samples) or (batch_size, 1, n_samples)"
+
+        if x.ndim == 3:
+            x = x.squeeze(1)
 
         x = self.model.forward(x, aug=False)
         if self.normalise_features:

oml/registry/models.py

@@ -4,6 +4,7 @@
 from torch import nn
 
 from oml.interfaces.models import IExtractor, IPairwiseModel
+from oml.models.audio.ecapa_tdnn.extractor import ECAPATDNNExtractor
 from oml.models.meta.projection import ExtractorWithMLP
 from oml.models.meta.siamese import (
     ConcatSiamese,
@@ -22,6 +23,7 @@
     "vit_clip": ViTCLIPExtractor,
     "vit_unicom": ViTUnicomExtractor,
     "extractor_with_mlp": ExtractorWithMLP,
+    "ecapa_tdnn": ECAPATDNNExtractor,
 }
 
 PAIRWISE_MODELS_REGISTRY = {