Merge pull request #35 from mhrice/cjs--classifier-v2

README.md

@@ -53,3 +53,15 @@ Apply remove effects: ['distortion'] (Up to 4, chosen randomly) -> Wet
 
 ## Misc.
 By default, files are rendered to `input_dir / processed / {string_of_effects} / {train|val|test}`.
+
+
+Download datasets:
+
+```
+python scripts/download.py vocalset guitarset idmt-smt-guitar idmt-smt-bass idmt-smt-drums
+```
+
+To run audio effects classifiction:
+```
+python scripts/train.py model=classifier "effects_to_use=[compressor, distortion, reverb, chorus, delay]" "effects_to_remove=[]" max_kept_effects=5 max_removed_effects=0 shuffle_kept_effects=True shuffle_removed_effects=True accelerator='gpu' render_root=/scratch/RemFX render_files=True
+```

remfx/datasets.py

@@ -18,9 +18,10 @@ from remfx.utils import create_sequential_chunks
 # https://zenodo.org/record/1193957 -> VocalSet
 
 ALL_EFFECTS = effects.Pedalboard_Effects
+print(ALL_EFFECTS)
 
 
-singer_splits = {
+vocalset_splits = {
     "train": [
         "male1",
         "male2",
@@ -43,6 +44,94 @@ singer_splits = {
     "test": ["male11", "female9"],
 }
 
+guitarset_splits = {"train": ["00", "01", "02", "03"], "val": ["04"], "test": ["05"]}
+idmt_guitar_splits = {
+    "train": ["classical", "country_folk", "jazz", "latin", "metal", "pop"],
+    "val": ["reggae", "ska"],
+    "test": ["rock", "blues"],
+}
+idmt_bass_splits = {
+    "train": ["BE", "BEQ"],
+    "val": ["VIF"],
+    "test": ["VIS"],
+}
+idmt_drums_splits = {
+    "train": ["WaveDrum02", "TechnoDrum01"],
+    "val": ["RealDrum01"],
+    "test": ["TechnoDrum02", "WaveDrum01"],
+}
+
+
+def locate_files(root: str, mode: str):
+    file_list = []
+    # ------------------------- VocalSet -------------------------
+    vocalset_dir = os.path.join(root, "VocalSet1-2")
+    if os.path.isdir(vocalset_dir):
+        # find all singer directories
+        singer_dirs = glob.glob(os.path.join(vocalset_dir, "data_by_singer", "*"))
+        singer_dirs = [
+            sd for sd in singer_dirs if os.path.basename(sd) in vocalset_splits[mode]
+        ]
+        files = []
+        for singer_dir in singer_dirs:
+            files += glob.glob(os.path.join(singer_dir, "**", "**", "*.wav"))
+        print(f"Found {len(files)} files in VocalSet {mode}.")
+        file_list += sorted(files)
+    # ------------------------- GuitarSet -------------------------
+    guitarset_dir = os.path.join(root, "audio_mono-mic")
+    if os.path.isdir(guitarset_dir):
+        files = glob.glob(os.path.join(guitarset_dir, "*.wav"))
+        files = [
+            f
+            for f in files
+            if os.path.basename(f).split("_")[0] in guitarset_splits[mode]
+        ]
+        print(f"Found {len(files)} files in GuitarSet {mode}.")
+        file_list += sorted(files)
+    # ------------------------- IDMT-SMT-GUITAR -------------------------
+    idmt_smt_guitar_dir = os.path.join(root, "IDMT-SMT-GUITAR_V2")
+    if os.path.isdir(idmt_smt_guitar_dir):
+        files = glob.glob(
+            os.path.join(
+                idmt_smt_guitar_dir, "IDMT-SMT-GUITAR_V2", "dataset4", "**", "*.wav"
+            ),
+            recursive=True,
+        )
+        files = [
+            f
+            for f in files
+            if os.path.basename(f).split("_")[0] in idmt_guitar_splits[mode]
+        ]
+        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):
+        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}.")
+    # ------------------------- IDMT-SMT-DRUMS -------------------------
+    idmt_smt_drums_dir = os.path.join(root, "IDMT-SMT-DRUMS-V2")
+    if os.path.isdir(idmt_smt_drums_dir):
+        files = glob.glob(os.path.join(idmt_smt_drums_dir, "audio", "*.wav"))
+        files = [
+            f
+            for f in files
+            if os.path.basename(f).split("_")[0] in idmt_drums_splits[mode]
+        ]
+        file_list += sorted(files)
+        print(f"Found {len(files)} files in IDMT-SMT-Drums {mode}.")
+
+    return file_list
+
 
 class VocalSet(Dataset):
     def __init__(
@@ -81,15 +170,7 @@ class VocalSet(Dataset):
         self.effects_to_keep = self.validate_effect_input()
         self.proc_root = self.render_root / "processed" / effects_string / self.mode
 
-        # find all singer directories
-        singer_dirs = glob.glob(os.path.join(self.root, "data_by_singer", "*"))
-        singer_dirs = [
-            sd for sd in singer_dirs if os.path.basename(sd) in singer_splits[mode]
-        ]
-        self.files = []
-        for singer_dir in singer_dirs:
-            self.files += glob.glob(os.path.join(singer_dir, "**", "**", "*.wav"))
-        self.files = sorted(self.files)
+        self.files = locate_files(self.root, self.mode)
 
         if self.proc_root.exists() and len(list(self.proc_root.iterdir())) > 0:
             print("Found processed files.")

scripts/download.py

@@ -0,0 +1,58 @@
+import os
+import sys
+import glob
+import torch
+import argparse
+
+
+def download_zip_dataset(dataset_url: str, output_dir: str):
+    zip_filename = os.path.basename(dataset_url)
+    zip_name = zip_filename.replace(".zip", "")
+    os.system(f"wget -P {output_dir} {dataset_url}")
+    os.system(
+        f"""unzip {os.path.join(output_dir, zip_filename)} -d {os.path.join(output_dir, zip_name)}"""
+    )
+    os.system(f"rm {os.path.join(output_dir, zip_filename)}")
+
+
+def process_dataset(dataset_dir: str, output_dir: str):
+    if dataset_dir == "VocalSet1-2":
+        pass
+    elif dataset_dir == "audio_mono-mic":
+        pass
+    elif dataset_dir == "IDMT-SMT-GUITAR_V2":
+        pass
+    elif dataset_dir == "IDMT-SMT-BASS":
+        pass
+    elif dataset_dir == "IDMT-SMT-DRUMS-V2":
+        pass
+    else:
+        raise NotImplemented(f"Invalid dataset_dir = {dataset_dir}.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "dataset_names",
+        choices=[
+            "vocalset",
+            "guitarset",
+            "idmt-smt-guitar",
+            "idmt-smt-bass",
+            "idmt-smt-drums",
+        ],
+        nargs="+",
+    )
+    args = parser.parse_args()
+
+    dataset_urls = {
+        "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",
+        "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")