feat: add support for Segment in one_hot_encoding

epratheeban · Sep 24, 2021 · 479e7b8 · 479e7b8
1 parent 5de2565
commit 479e7b8
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/pyannote/core/utils/numpy.py b/pyannote/core/utils/numpy.py
@@ -41,7 +41,7 @@
 
 def one_hot_encoding(
     annotation: Annotation,
-    support: Timeline,
+    support: Union[Segment, Timeline],
     window: Union[SlidingWindow, SlidingWindowFeature],
     labels: Optional[List[Text]] = None,
     mode: Text = "center",
@@ -51,7 +51,7 @@ def one_hot_encoding(
     Parameters
     ----------
     annotation : `pyannote.core.Annotation`
-    support : `pyannote.core.Timeline`
+    support : `pyannote.core.Timeline` or `pyannote.core.Segment`
     window : `SlidingWindow`
         Use this `window`.
     labels : list, optional
@@ -79,7 +79,12 @@ def one_hot_encoding(
             )
             raise TypeError(msg)
 
-    extent = support.extent()
+    if isinstance(support, Segment):
+        extent = support
+        support = Timeline([support])
+    else:
+        extent = support.extent()
+
     window = SlidingWindow(
         start=extent.start, step=window.step, duration=window.duration
     )