ycemsubakan · Sep 5, 2022 · Jul 13, 2022 · Jul 13, 2022 · Jul 13, 2022 · Jul 31, 2022
diff --git a/recipes/BinauralWSJ0Mix/prepare_data.py b/recipes/BinauralWSJ0Mix/prepare_data.py
@@ -3,7 +3,7 @@
 
 Author
  * Cem Subakan 2020
- * Zijian
+ * Zijian 2022
 
  """
 
@@ -12,7 +12,13 @@
 
 
 def prepare_binaural_wsj0mix(
-    datapath, savepath, n_spks=2, skip_prep=False, fs=8000, version="min",
+    experiment_name,
+    datapath,
+    savepath,
+    n_spks=2,
+    skip_prep=False,
+    fs=8000,
+    version="min",
 ):
     """
     Prepared binaural wsj2mix if n_spks=2 and binaural wsj3mix if n_spks=3.
@@ -28,12 +34,16 @@ def prepare_binaural_wsj0mix(
     if skip_prep:
         return
 
-    if n_spks == 2:
+    if "noise" in experiment_name:
+        create_binaural_wsj0mix2_noise_csv(datapath, savepath, fs, version)
+    elif "reverb" in experiment_name:
+        create_binaural_wsj0mix2_reverb_csv(datapath, savepath, fs, version)
+    elif n_spks == 2:
         create_binaural_wsj0mix2_csv(datapath, savepath, fs, version)
     elif n_spks == 3:
         create_binaural_wsj0mix3_csv(datapath, savepath, fs, version)
     else:
-        raise ValueError("Unsupported Number of Speakers")
+        raise ValueError("Unsupported preparation mode")
 
 
 def create_binaural_wsj0mix2_csv(
@@ -202,3 +212,179 @@ def create_binaural_wsj0mix3_csv(
                     "s3_wav_opts": None,
                 }
                 writer.writerow(row)
+
+
+def create_binaural_wsj0mix2_noise_csv(
+    datapath,
+    savepath,
+    fs,
+    version,
+    savename="binaural_wsj0-2mix_noise_",
+    set_types=["tr", "cv", "tt"],
+):
+    """
+    This function creates the csv files to get the speechbrain data loaders for the whamr dataset.
+
+    Arguments:
+        datapath (str) : path for the wsj0-mix dataset.
+        savepath (str) : path where we save the csv file
+        fs (int) : the sampling rate
+        version (str) : min or max
+        savename (str) : the prefix to use for the .csv files
+        set_types (list) : the sets to create
+    """
+    if fs == 8000:
+        sample_rate = "8k"
+    elif fs == 16000:
+        sample_rate = "16k"
+    else:
+        raise ValueError("Unsupported sampling rate")
+
+    for set_type in set_types:
+
+        mix_path = os.path.join(
+            datapath,
+            "wav{}".format(sample_rate),
+            version,
+            set_type,
+            "mix_both/",
+        )
+        s1_path = os.path.join(
+            datapath, "wav{}".format(sample_rate), version, set_type, "s1/",
+        )
+        s2_path = os.path.join(
+            datapath, "wav{}".format(sample_rate), version, set_type, "s2/",
+        )
+        noise_path = os.path.join(
+            datapath, "wav{}".format(sample_rate), version, set_type, "noise/"
+        )
+
+        files = os.listdir(mix_path)
+
+        mix_fl_paths = [mix_path + fl for fl in files]
+        s1_fl_paths = [s1_path + fl for fl in files]
+        s2_fl_paths = [s2_path + fl for fl in files]
+        noise_fl_paths = [noise_path + fl for fl in files]
+
+        csv_columns = [
+            "ID",
+            "duration",
+            "mix_wav",
+            "mix_wav_format",
+            "mix_wav_opts",
+            "s1_wav",
+            "s1_wav_format",
+            "s1_wav_opts",
+            "s2_wav",
+            "s2_wav_format",
+            "s2_wav_opts",
+            "noise_wav",
+            "noise_wav_format",
+            "noise_wav_opts",
+        ]
+
+        with open(
+            os.path.join(savepath, savename + set_type + ".csv"), "w"
+        ) as csvfile:
+            writer = csv.DictWriter(csvfile, fieldnames=csv_columns)
+            writer.writeheader()
+            for (i, (mix_path, s1_path, s2_path, noise_path),) in enumerate(
+                zip(mix_fl_paths, s1_fl_paths, s2_fl_paths, noise_fl_paths,)
+            ):
+
+                row = {
+                    "ID": i,
+                    "duration": 1.0,
+                    "mix_wav": mix_path,
+                    "mix_wav_format": "wav",
+                    "mix_wav_opts": None,
+                    "s1_wav": s1_path,
+                    "s1_wav_format": "wav",
+                    "s1_wav_opts": None,
+                    "s2_wav": s2_path,
+                    "s2_wav_format": "wav",
+                    "s2_wav_opts": None,
+                    "noise_wav": noise_path,
+                    "noise_wav_format": "wav",
+                    "noise_wav_opts": None,
+                }
+                writer.writerow(row)
+
+
+def create_binaural_wsj0mix2_reverb_csv(
+    datapath,
+    savepath,
+    fs,
+    version,
+    savename="binaural_wsj0-2mix_reverb_",
+    set_types=["tr", "cv", "tt"],
+):
+    """
+    This function creates the csv files to get the speechbrain data loaders for the binaural-wsj0mix2 dataset.
+
+    Arguments:
+        datapath (str) : path for the binaural-wsj0mix dataset.
+        savepath (str) : path where we save the csv file
+    """
+    if fs == 8000:
+        sample_rate = "8k"
+    elif fs == 16000:
+        sample_rate = "16k"
+    else:
+        raise ValueError("Unsupported sampling rate")
+
+    for set_type in set_types:
+
+        mix_path = os.path.join(
+            datapath, "wav{}".format(sample_rate), version, set_type, "mix/",
+        )
+        s1_path = os.path.join(
+            datapath, "wav{}".format(sample_rate), version, set_type, "s1/",
+        )
+        s2_path = os.path.join(
+            datapath, "wav{}".format(sample_rate), version, set_type, "s2/",
+        )
+
+        files = os.listdir(mix_path)
+
+        mix_fl_paths = [mix_path + fl for fl in files]
+        s1_fl_paths = [s1_path + fl for fl in files]
+        s2_fl_paths = [s2_path + fl for fl in files]
+
+        csv_columns = [
+            "ID",
+            "duration",
+            "mix_wav",
+            "mix_wav_format",
+            "mix_wav_opts",
+            "s1_wav",
+            "s1_wav_format",
+            "s1_wav_opts",
+            "s2_wav",
+            "s2_wav_format",
+            "s2_wav_opts",
+        ]
+
+        with open(
+            os.path.join(savepath, savename + set_type + ".csv"), "w"
+        ) as csvfile:
+            writer = csv.DictWriter(csvfile, fieldnames=csv_columns)
+            writer.writeheader()
+            for i, (mix_path, s1_path, s2_path) in enumerate(
+                zip(mix_fl_paths, s1_fl_paths, s2_fl_paths)
+            ):
+
+                row = {
+                    "ID": i,
+                    "duration": 1.0,
+                    "mix_wav": mix_path,
+                    "mix_wav_format": "wav",
+                    "mix_wav_opts": None,
+                    "s1_wav": s1_path,
+                    "s1_wav_format": "wav",
+                    "s1_wav_opts": None,
+                    "s2_wav": s2_path,
+                    "s2_wav_format": "wav",
+                    "s2_wav_opts": None,
+                }
+                writer.writerow(row)
diff --git a/recipes/BinauralWSJ0Mix/separation/README.md b/recipes/BinauralWSJ0Mix/separation/README.md
@@ -15,13 +15,13 @@ python train.py hparams/convtasnet-parallel.yaml
                 --data_folder yourpath/binaural-wsj0mix/2speakers
                 --wsj_root yourpath/to/wsj/
 ```
-The data_folder will be automatically created from the wsj_root one.
+The training data will be automatically created from the `wsj_root`, which is the root folder that contains
 Note that during training we print the negative SNR instead of SI-SNR because the scale-invariance property of SI-SNR makes it insensitive to power rescaling of the estimated signal, which may fail in preserving the ILD between the outputs.
 
 
 # Binaural WSJ0-2mix and WSJ0-3mix dataset creation
-* The best way to create the datasets is using the the scripts at https://github.com/huangzj421/Binaural-WSJ0Mix. The train.py will download it under data_folder if not exists.
-* Also the train.py will create Binaural WSJ0Mix datasets automatically with wsj_root specified:`python train.py hparams/convtasnet-parallel.yaml --data_folder yourpath/binaural-wsj0mix/2speakers --wsj_root yourpath/to/wsj/`
+* The training data generation scripts can be found from [https://github.com/huangzj421/Binaural-WSJ0Mix](https://github.com/huangzj421/Binaural-WSJ0Mix). But the `train.py` also automatically downloads and generates the data. It puts the data under the path specified in `data_folder`.
+* The default command to run that automatically generate the data given wsj0 folder:`python train.py hparams/convtasnet-parallel.yaml --data_folder yourpath/binaural-wsj0mix/2speakers --wsj_root yourpath/wsj0-mix/wsj0`
 
 
 # Dynamic Mixing:
@@ -37,11 +37,15 @@ Here are the SNRi results (in dB) as well as ITD and ILD errors as the metric fo
 | --- | --- | --- | --- |
 |ConvTasnet-independent| 12.39 | 6.17 | 0.32 |
 |ConvTasnet-cross| 11.9 | 5.69 | 0.37 |
-|ConvTasnet-parallel| 11.69 | 3.23 | 0.23 |
+|ConvTasnet-parallel| 16.93 | 2.38 | 0.09 |
+|ConvTasnet-parallel-noise| 18.25 | 5.56 | 0.23 |
+|ConvTasnet-parallel-reverb| 6.95 | 0.0 | 0.40 |
 
 * ConvTasnet-independent.yaml refers to ConvTasnet is applied to each channel independently.
 * ConvTasnet-cross.yaml refers to cross-channel features like ILD or IPD are concatenated to the encoder output.
 * ConvTasnet-parallel.yaml refers to the proposed multiinput-multi-output (MIMO) TasNet in the original paper [1].
+* ConvTasnet-parallel-noise.yaml refers to the above Tasnet applied to 2 speakers with DEMAND noise.
+* ConvTasnet-parallel-reverb.yaml refers to the above Tasnet applied to 2 speakers with reverberance(RT60) from the [BRIR Sim Set](http://iosr.uk/software/index.php).
 
 The output folders with the checkpoints, logs, etc are available [here](https://drive.google.com/drive/folders/17FFwlIq6MQLHT9RXPgeYssti5TEeEXsx?usp=sharing)
 

diff --git a/recipes/BinauralWSJ0Mix/separation/dynamic_mixing.py b/recipes/BinauralWSJ0Mix/separation/dynamic_mixing.py
@@ -4,7 +4,6 @@
 import torchaudio
 import glob
 import os
-from pathlib import Path
 import random
 from speechbrain.processing.signal_processing import rescale
 from speechbrain.dataio.batch import PaddedBatch
@@ -43,7 +42,7 @@ def dynamic_mix_data_prep(hparams):
     spk_list = [x for x in spk_hashtable.keys()]
     spk_weights = [x / sum(spk_weights) for x in spk_weights]
 
-    if "wham" in Path(hparams["data_folder"]).stem:
+    if "noise" in hparams["experiment_name"]:
         noise_files = get_wham_noise_filenames(hparams)
 
     @sb.utils.data_pipeline.takes("mix_wav")
@@ -61,7 +60,7 @@ def audio_pipeline(
             spk_list, hparams["num_spks"], replace=False, p=spk_weights
         )
 
-        if "wham" in Path(hparams["data_folder"]).stem:
+        if "noise" in hparams["experiment_name"]:
             noise_file = np.random.choice(noise_files, 1, replace=False)
 
             noise, fs_read = torchaudio.load(noise_file[0])
@@ -108,30 +107,48 @@ def audio_pipeline(
                 )
                 tmp = rescale(tmp, torch.tensor(len(tmp)), gain, scale="dB")
 
-            tmp_bi = torch.FloatTensor(len(tmp), 2)  # binaural
-            subject_path_list = glob.glob(
-                os.path.join(hparams["hrtf_wav_path"], "subject*")
-            )
-            subject_path = np.random.choice(subject_path_list)
-            azimuth_list = (
-                [-80, -65, -55] + list(range(-45, 46, 5)) + [55, 65, 80]
-            )
-            azimuth = np.random.choice(azimuth_list)
-
-            for i, loc in enumerate(["left", "right"]):
-
+            if "reverb" in hparams["experiment_name"]:
+                tmp = torch.stack((tmp, tmp), 1)
+                reverb_time = np.random.choice(
+                    ["0_1s", "0_2s", "0_4s", "0_7s", "0_8s"]
+                )
+                azimuth = np.random.choice(list(range(-90, 91, 5)))
                 hrtf_file = os.path.join(
-                    subject_path,
-                    "{}az{}.wav".format(
-                        azimuth.astype("str").replace("-", "neg"), loc
-                    ),
+                    hparams["hrtf_wav_path"],
+                    reverb_time,
+                    "CATT_{}_{}.wav".format(reverb_time, azimuth),
                 )
                 hrtf, sr = torchaudio.load(hrtf_file)
                 transform = torchaudio.transforms.Resample(sr, fs_read)
-                hrtf = transform(hrtf[:, np.random.randint(50)])
-                tmp_bi[:, i] = torch.from_numpy(
+                hrtf = transform(hrtf)
+                tmp_bi = torch.from_numpy(
                     fftconvolve(tmp.numpy(), hrtf.numpy(), mode="same")
                 )
+            else:
+                tmp_bi = torch.FloatTensor(len(tmp), 2)  # binaural
+                subject_path_list = glob.glob(
+                    os.path.join(hparams["hrtf_wav_path"], "subject*")
+                )
+                subject_path = np.random.choice(subject_path_list)
+                azimuth_list = (
+                    [-80, -65, -55] + list(range(-45, 46, 5)) + [55, 65, 80]
+                )
+                azimuth = np.random.choice(azimuth_list)
+
+                for i, loc in enumerate(["left", "right"]):
+
+                    hrtf_file = os.path.join(
+                        subject_path,
+                        "{}az{}.wav".format(
+                            azimuth.astype("str").replace("-", "neg"), loc
+                        ),
+                    )
+                    hrtf, sr = torchaudio.load(hrtf_file)
+                    transform = torchaudio.transforms.Resample(sr, fs_read)
+                    hrtf = transform(hrtf[:, np.random.randint(50)])
+                    tmp_bi[:, i] = torch.from_numpy(
+                        fftconvolve(tmp.numpy(), hrtf.numpy(), mode="same")
+                    )
 
             # Make relative source energy same with original
             spatial_scaling = torch.sqrt(
@@ -149,10 +166,11 @@ def audio_pipeline(
 
         sources = torch.stack(sources)
         mixture = torch.sum(sources, 0)
-        if "wham" in Path(hparams["data_folder"]).stem:
+        if "noise" in hparams["experiment_name"]:
             len_noise = len(noise)
             len_mix = len(mixture)
             min_len = min(len_noise, len_mix)
+            noise = torch.swapaxes(noise, 0, 1)
             mixture = mixture[:min_len] + noise[:min_len]
 
         max_amp = max(
@@ -174,7 +192,7 @@ def audio_pipeline(
         if hparams["num_spks"] == 2:
             yield None
 
-        if "wham" in Path(hparams["data_folder"]).stem:
+        if "noise" in hparams["experiment_name"]:
             mean_source_lvl = sources.abs().mean()
             mean_noise_lvl = noise.abs().mean()
             noise = (mean_source_lvl / mean_noise_lvl) * noise