Make Compose just like Sequential.

descriptinc · Mar 10, 2022 · 1a7a186 · 1a7a186
1 parent 63f99ea
commit 1a7a186
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 7 deletions.
diff --git a/audiotools/data/transforms.py b/audiotools/data/transforms.py
@@ -115,12 +115,8 @@ def batch_instantiate(
 
 class Compose(BaseTransform):
     def __init__(self, transforms: list, name: str = None, prob: float = 1.0):
-        tfm_counts = defaultdict(lambda: 0)
-        for tfm in transforms:
-            tfm_name = tfm.name
-            tfm_name = f"{tfm_counts[tfm_name]}.{tfm_name}"
-            tfm_counts[tfm.name] += 1
-            tfm.name = tfm_name
+        for i, tfm in enumerate(transforms):
+            tfm.name = f"{i}.{tfm.name}"
 
         keys = [tfm.name for tfm in transforms]
         super().__init__(keys=keys, name=name, prob=prob)

diff --git a/tests/data/test_datasets.py b/tests/data/test_datasets.py
@@ -138,7 +138,7 @@ def test_csv_dataset():
         signal = dataset.transform(signal, **kwargs)
         original = dataset.transform(original, **kwargs)
 
-        mask = kwargs["Compose"]["0.Silence"]["mask"]
+        mask = kwargs["Compose"]["1.Silence"]["mask"]
 
         zeros_ = torch.zeros_like(signal[mask].audio_data)
         original_ = original[~mask].audio_data