Add dsd100 dataset

cfg/config.yaml

@@ -53,9 +53,9 @@ callbacks:
     _target_: remfx.callbacks.MetricCallback
 
 datamodule:
-  _target_: remfx.datasets.VocalSetDatamodule
+  _target_: remfx.datasets.EffectDatamodule
   train_dataset:
-    _target_: remfx.datasets.VocalSet
+    _target_: remfx.datasets.EffectDataset
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
     chunk_size: ${chunk_size}
@@ -70,7 +70,7 @@ datamodule:
     render_files: ${render_files}
     render_root: ${render_root}
   val_dataset:
-    _target_: remfx.datasets.VocalSet
+    _target_: remfx.datasets.EffectDataset
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
     chunk_size: ${chunk_size}
@@ -85,7 +85,7 @@ datamodule:
     render_files: ${render_files}
     render_root: ${render_root}
   test_dataset:
-    _target_: remfx.datasets.VocalSet
+    _target_: remfx.datasets.EffectDataset
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
     chunk_size: ${chunk_size}

cfg/exp/default.yaml

@@ -1,6 +1,6 @@
 # @package _global_
 defaults:
-  - override /model: demucs
+  - override /model: umx
   - override /effects: all
 seed: 12345
 sample_rate: 48000

remfx/datasets.py

@@ -55,6 +55,11 @@ idmt_bass_splits = {
     "val": ["VIF"],
     "test": ["VIS"],
 }
+dsd_100_splits = {
+    "train": ["train"],
+    "val": ["val"],
+    "test": ["test"],
+}
 idmt_drums_splits = {
     "train": ["WaveDrum02", "TechnoDrum01"],
     "val": ["RealDrum01"],
@@ -105,19 +110,28 @@ def locate_files(root: str, mode: str):
         file_list += sorted(files)
         print(f"Found {len(files)} files in IDMT-SMT-Guitar {mode}.")
     # ------------------------- IDMT-SMT-BASS -------------------------
-    idmt_smt_bass_dir = os.path.join(root, "IDMT-SMT-BASS")
-    if os.path.isdir(idmt_smt_bass_dir):
+    # idmt_smt_bass_dir = os.path.join(root, "IDMT-SMT-BASS")
+    # if os.path.isdir(idmt_smt_bass_dir):
+    #     files = glob.glob(
+    #         os.path.join(idmt_smt_bass_dir, "**", "*.wav"),
+    #         recursive=True,
+    #     )
+    #     files = [
+    #         f
+    #         for f in files
+    #         if os.path.basename(os.path.dirname(f)) in idmt_bass_splits[mode]
+    #     ]
+    #     file_list += sorted(files)
+    #     print(f"Found {len(files)} files in IDMT-SMT-Bass {mode}.")
+    # ------------------------- DSD100 ---------------------------------
+    dsd_100_dir = os.path.join(root, "DSD100")
+    if os.path.isdir(dsd_100_dir):
         files = glob.glob(
-            os.path.join(idmt_smt_bass_dir, "**", "*.wav"),
+            os.path.join(dsd_100_dir, mode, "**", "*.wav"),
             recursive=True,
         )
-        files = [
-            f
-            for f in files
-            if os.path.basename(os.path.dirname(f)) in idmt_bass_splits[mode]
-        ]
         file_list += sorted(files)
-        print(f"Found {len(files)} files in IDMT-SMT-Bass {mode}.")
+        print(f"Found {len(files)} files in DSD100 {mode}.")
     # ------------------------- IDMT-SMT-DRUMS -------------------------
     idmt_smt_drums_dir = os.path.join(root, "IDMT-SMT-DRUMS-V2")
     if os.path.isdir(idmt_smt_drums_dir):
@@ -133,7 +147,7 @@ def locate_files(root: str, mode: str):
     return file_list
 
 
-class VocalSet(Dataset):
+class EffectDataset(Dataset):
     def __init__(
         self,
         root: str,
@@ -199,6 +213,9 @@ class VocalSet(Dataset):
                     if resampled_chunk.shape[-1] < chunk_size:
                         # Skip if chunk is too small
                         continue
+                    # Sum to mono
+                    if resampled_chunk.shape[0] > 1:
+                        resampled_chunk = resampled_chunk.sum(0, keepdim=True)
 
                     dry, wet, dry_effects, wet_effects = self.process_effects(
                         resampled_chunk
@@ -334,7 +351,7 @@ class VocalSet(Dataset):
         return normalized_dry, normalized_wet, dry_labels_tensor, wet_labels_tensor
 
 
-class VocalSetDatamodule(pl.LightningDataModule):
+class EffectDatamodule(pl.LightningDataModule):
     def __init__(
         self,
         train_dataset,

remfx/models.py

@@ -2,7 +2,6 @@ import torch
 import torchmetrics
 import pytorch_lightning as pl
 from torch import Tensor, nn
-from torch.nn import functional as F
 from torchaudio.models import HDemucs
 from audio_diffusion_pytorch import DiffusionModel
 from auraloss.time import SISDRLoss

scripts/download.py

@@ -1,8 +1,6 @@
 import os
-import sys
-import glob
-import torch
 import argparse
+import shutil
 
 
 def download_zip_dataset(dataset_url: str, output_dir: str):
@@ -26,8 +24,42 @@ def process_dataset(dataset_dir: str, output_dir: str):
         pass
     elif dataset_dir == "IDMT-SMT-DRUMS-V2":
         pass
+    elif dataset_dir == "DSD100":
+        shutil.rmtree(os.path.join(output_dir, dataset_dir, "Mixtures"))
+        for dir in os.listdir(os.path.join(output_dir, dataset_dir, "Sources", "Dev")):
+            source = os.path.join(output_dir, dataset_dir, "Sources", "Dev", dir)
+            shutil.move(source, os.path.join(output_dir, dataset_dir))
+        shutil.rmtree(os.path.join(output_dir, dataset_dir, "Sources", "Dev"))
+        for dir in os.listdir(os.path.join(output_dir, dataset_dir, "Sources", "Test")):
+            source = os.path.join(output_dir, dataset_dir, "Sources", "Test", dir)
+            shutil.move(source, os.path.join(output_dir, dataset_dir))
+        shutil.rmtree(os.path.join(output_dir, dataset_dir, "Sources", "Test"))
+        shutil.rmtree(os.path.join(output_dir, dataset_dir, "Sources"))
+
+        os.mkdir(os.path.join(output_dir, dataset_dir, "train"))
+        os.mkdir(os.path.join(output_dir, dataset_dir, "val"))
+        os.mkdir(os.path.join(output_dir, dataset_dir, "test"))
+        files = os.listdir(os.path.join(output_dir, dataset_dir))
+
+        num = 0
+        for dir in files:
+            if not os.path.isdir(os.path.join(output_dir, dataset_dir, dir)):
+                continue
+            if dir == "train" or dir == "val" or dir == "test":
+                continue
+            source = os.path.join(output_dir, dataset_dir, dir, "bass.wav")
+            if num < 80:
+                dest = os.path.join(output_dir, dataset_dir, "train", f"{num}.wav")
+            elif num < 90:
+                dest = os.path.join(output_dir, dataset_dir, "val", f"{num}.wav")
+            else:
+                dest = os.path.join(output_dir, dataset_dir, "test", f"{num}.wav")
+            shutil.move(source, dest)
+            shutil.rmtree(os.path.join(output_dir, dataset_dir, dir))
+            num += 1
+
     else:
-        raise NotImplemented(f"Invalid dataset_dir = {dataset_dir}.")
+        raise NotImplementedError(f"Invalid dataset_dir = {dataset_dir}.")
 
 
 if __name__ == "__main__":
@@ -38,7 +70,7 @@ if __name__ == "__main__":
             "vocalset",
             "guitarset",
             "idmt-smt-guitar",
-            "idmt-smt-bass",
+            "dsd100",
             "idmt-smt-drums",
         ],
         nargs="+",
@@ -49,10 +81,11 @@ if __name__ == "__main__":
         "vocalset": "https://zenodo.org/record/1442513/files/VocalSet1-2.zip",
         "guitarset": "https://zenodo.org/record/3371780/files/audio_mono-mic.zip",
         "IDMT-SMT-GUITAR_V2": "https://zenodo.org/record/7544110/files/IDMT-SMT-GUITAR_V2.zip",
-        "IDMT-SMT-BASS": "https://zenodo.org/record/7188892/files/IDMT-SMT-BASS.zip",
+        "DSD100": "http://liutkus.net/DSD100.zip",
         "IDMT-SMT-DRUMS-V2": "https://zenodo.org/record/7544164/files/IDMT-SMT-DRUMS-V2.zip",
     }
 
     for dataset_name, dataset_url in dataset_urls.items():
         if dataset_name in args.dataset_names:
             download_zip_dataset(dataset_url, "~/data/remfx-data")
+            process_dataset(dataset_name, "~/data/remfx-data")

shell_vars.sh

@@ -1,3 +1,3 @@
-export DATASET_ROOT="./data/VocalSet"
+export DATASET_ROOT="./data/"
 export WANDB_PROJECT="RemFX"
 export WANDB_ENTITY="mattricesound"