More cleanup, edit configs

README.md

@@ -1,6 +1,10 @@
 # General Purpose Audio Effect Removal
-# About
-TBD. Add photo. Add paper link.
+Removing multiple audio effects from multiple sources using compositional audio effect removal and source separation and speech enhancement models.
+
+This repo contains the code for the paper [General Purpose Audio Effect Removal](https://arxiv.org/abs/2110.00484). (Todo: Link broken, Add video, Add img)
+
+
+
 # Setup
 ```
 git clone https://github.com/mhrice/RemFx.git
@@ -8,6 +12,7 @@ git submodule update --init --recursive
 pip install . umx
 ```
 # Usage
+This repo can be used for many different tasks. Here are some examples.
 ## Run RemFX Detect on a single file
 ```
 ./download_checkpoints.sh
@@ -21,24 +26,25 @@ unzip RemFX_eval_dataset.zip
 
 ## Download the datasets used in the paper
 ```
-python scripts/download.py vocalset guitarset idmt-smt-guitar idmt-smt-bass idmt-smt-drums
+python scripts/download.py vocalset guitarset idmt-smt-bass idmt-smt-drums
 ```
+By default, the datasets are downloaded to `./data/remfx-data`. To change this, pass `--output_dir={path/to/datasets}` to `download.py`
 
-
-## Training
-Before training, it is important that you have downloaded the datasets (see above).
-This project uses [hydra](https://hydra.cc/) for configuration management. All experiments are defined in `cfg/exp/`. To train with an existing experiment, first run
+Then set the dataset root :
 ```
 export DATASET_ROOT={path/to/datasets}
 ```
-Then:
+
+## Training
+Before training, it is important that you have downloaded the datasets (see above) and set DATASET_ROOT.
+This project uses the [pytorch-lightning](https://www.pytorchlightning.ai/index.html) framework and [hydra](https://hydra.cc/) for configuration management. All experiments are defined in `cfg/exp/`. To train with an existing experiment run
 ```
 python scripts/train.py +exp={experiment_name}
 ```
 
 Here are some selected experiment types from the paper, which use different datasets and configurations. See `cfg/exp/` for a full list of experiments and parameters.
 
-| Experiment Type         | config name  | example          |
+| Experiment Type         | Config Name  | Example          |
 | ----------------------- | ------------ | ---------------- |
 | Effect-specific         | {effect}     | +exp=chorus      |
 | Effect-specific + FXAug | {effect}_aug | +exp=chorus_aug  |
@@ -49,6 +55,16 @@ Here are some selected experiment types from the paper, which use different data
 To change the configuration, simply edit the experiment file, or override the configuration on the command line. A description of some of these variables is in the Misc. section below.
 You can also create a custom experiment by creating a new experiment file in `cfg/exp/` and overriding the default parameters in `config.yaml`.
 
+At the end of training, the train script will automatically evaluate the test set using the best checkpoint (by validation loss). To evaluate a specific checkpoint, run
+
+```
+python test.py +exp={experiment_name} ckpt_path={path/to/checkpoint}
+```
+
+If you have generated the dataset separately from training, be sure to set `render_files=False` in the config or command-line, and set `render_root={path_to_dataset}` if it is in a custom location.
+
+Also note that the training assumes you have a GPU. To train on CPU, set `accelerator=null` in the config or command-line.
+
 ## Evaluate models on the General Purpose Audio Effect Removal evaluation dataset
 First download the dataset (see above).
 To use the pretrained RemFX model, download the checkpoints
@@ -70,9 +86,15 @@ Download checkpoints from [here](https://zenodo.org/record/8179396), or see the
 
 
 ## Generate datasets used in the paper
+Before generating datasets, it is important that you have downloaded the datasets (see above) and set DATASET_ROOT.
+
+To generate one of the datasets used in the paper, it is as simple as running a training job with a particular config. For example, to generate the `chorus` FXAug dataset, which includes files with 5 possible effects, up to 4 kept effects (distortion, reverb, compression, delay), and 1 removed effects (chorus), run
 ```
+python scripts/train.py +exp=chorus_aug
 ```
-Note that by default, files are rendered to `input_dir / processed / {string_of_effects} / {train|val|test}`.
+
+See the Misc. section below for a description of the parameters.
+By default, files are rendered to `{render_root} / processed / {string_of_effects} / {train|val|test}`.
 
 ## Evaluate with a custom directory
 Assumes directory is structured as
@@ -86,21 +108,27 @@ Assumes directory is structured as
         - file2.wav
         - file3.wav
 
-Change root path in `shell_vars.sh` and `source shell_vars.sh`
+First set the dataset root:
+```
+export DATASET_ROOT={path/to/datasets}
+```
 
-`python scripts/chain_inference.py +exp=chain_inference_custom`
+Then run
+```
+python scripts/chain_inference.py +exp=chain_inference_custom
+```
 
 # Misc.
 ## Experimental parameters
 Some relevant training parameters descriptions
 - `num_kept_effects={[min, max]}` range of <b> Kept </b> effects to apply to each file. Inclusive.
 - `num_removed_effects={[min, max]}` range of <b> Removed </b> effects to apply to each file. Inclusive.
-- `model={model}` architecture to use (see 'Models')
+- `model={model}` architecture to use (see 'Effect Removal Models/Effect Classification Models')
 - `effects_to_keep={[effect]}` Effects to apply but not remove (see 'Effects')
 - `effects_to_remove={[effect]}` Effects to remove (see 'Effects')
 - `accelerator=null/'gpu'` Use GPU (1 device) (default: null)
 - `render_files=True/False` Render files. Disable to skip rendering stage (default: True)
-- `render_root={path/to/dir}`. Root directory to render files to (default: DATASET_ROOT)
+- `render_root={path/to/dir}`. Root directory to render files to (default: ./data)
 
 ### Effect Removal Models
 - `umx`
@@ -121,10 +149,3 @@ Some relevant training parameters descriptions
 - `distortion`
 - `reverb`
 - `delay`
-
-
-
-
-
-
-

cfg/exp/1-1.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/2-2.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/3-3.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/4-4.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/5-1.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/5-5.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/5-5_cls.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "/scratch/cjs-logs"
 render_files: True
-render_root: "/scratch/EffectSet_cjs"
+
 accelerator: "gpu"
 log_audio: False
 # Effects

cfg/exp/5-5_cls_dynamic.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "/scratch/cjs-logs"
 render_files: True
-render_root: "/scratch/EffectSet_cjs"
+
 accelerator: "gpu"
 log_audio: False
 # Effects

cfg/exp/chain_inference.yaml

@@ -6,7 +6,7 @@ seed: 12345
 sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/chain_inference_aug.yaml

@@ -6,7 +6,7 @@ seed: 12345
 sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/chain_inference_aug_classifier.yaml

@@ -6,7 +6,7 @@ seed: 12345
 sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/chain_inference_custom.yaml

@@ -6,7 +6,7 @@ seed: 12345
 sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/chorus.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/chorus_aug.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/compression.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/compression_aug.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/delay.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/delay_aug.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/distortion.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/distortion_aug.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/reverb.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

cfg/exp/reverb_aug.yaml

@@ -7,7 +7,7 @@ sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "/scratch/EffectSet"
+
 accelerator: "gpu"
 log_audio: True
 # Effects

remfx/datasets.py

@@ -44,16 +44,6 @@ vocalset_splits = {
 }
 
 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"],
-}
 dsd_100_splits = {
     "train": ["train"],
     "val": ["val"],
@@ -92,36 +82,6 @@ def locate_files(root: str, mode: str):
         ]
         print(f"Found {len(files)} files in GuitarSet {mode}.")
         file_list.append(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.append(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.append(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):

remfx/models.py

@@ -9,11 +9,9 @@ from auraloss.time import SISDRLoss
 from auraloss.freq import MultiResolutionSTFTLoss
 from umx.openunmix.model import OpenUnmix, Separator
 
-from remfx.utils import FADLoss, spectrogram
+from remfx.utils import spectrogram
 from remfx.tcn import TCN
 from remfx.utils import causal_crop
-from remfx.callbacks import log_wandb_audio_batch
-from einops import rearrange
 from remfx import effects
 import asteroid
 import random
@@ -148,13 +146,14 @@ class RemFXChainInference(pl.LightningModule):
                 )
                 # print(f"Input_{metric}", negate * self.metrics[metric](x, y))
                 # print(f"test_{metric}", negate * self.metrics[metric](output, y))
-                self.output_str += f"{negate * self.metrics[metric](x, y).item():.4f},{negate * self.metrics[metric](output, y).item():.4f},"
-            self.output_str += "\n"
+                # self.output_str += f"{negate * self.metrics[metric](x, y).item():.4f},{negate * self.metrics[metric](output, y).item():.4f},"
+            # self.output_str += "\n"
         return loss
 
     def on_test_end(self) -> None:
-        with open("output.csv", "w") as f:
-            f.write(self.output_str)
+        pass
+        # with open("output.csv", "w") as f:
+        # f.write(self.output_str)
 
     def sample(self, batch):
         return self.forward(batch, 0)[1]
@@ -266,13 +265,14 @@ class RemFX(pl.LightningModule):
                 )
                 # print(f"Input_{metric}", negate * self.metrics[metric](x, y))
                 # print(f"test_{metric}", negate * self.metrics[metric](output, y))
-                self.output_str += f"{negate * self.metrics[metric](x, y).item():.4f},{negate * self.metrics[metric](output, y).item():.4f},"
-            self.output_str += "\n"
+                # self.output_str += f"{negate * self.metrics[metric](x, y).item():.4f},{negate * self.metrics[metric](output, y).item():.4f},"
+            # self.output_str += "\n"
         return loss
 
     def on_test_end(self) -> None:
-        with open("output.csv", "w") as f:
-            f.write(self.output_str)
+        pass
+        # with open("output.csv", "w") as f:
+        # f.write(self.output_str)
 
 
 class OpenUnmixModel(nn.Module):

remfx/tcn.py

@@ -125,7 +125,6 @@ class TCN(nn.Module):
         self.buffer = torch.zeros(2, self.receptive_field + self.block_size - 1)
 
     def forward(self, x: Tensor) -> Tensor:
-        x_in = x
         for _, block in enumerate(self.process_blocks):
             x = block(x)
         y_hat = torch.tanh(self.output(x))

scripts/download.py

@@ -18,8 +18,6 @@ def process_dataset(dataset_dir: str, output_dir: str):
         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":
@@ -69,23 +67,25 @@ if __name__ == "__main__":
         choices=[
             "vocalset",
             "guitarset",
-            "idmt-smt-guitar",
             "dsd100",
             "idmt-smt-drums",
         ],
         nargs="+",
     )
+    parser.add_argument("--output_dir", default="./data/remfx-data")
     args = parser.parse_args()
 
+    if not os.path.exists(args.output_dir):
+        os.makedirs(args.output_dir)
+
     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",
         "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")
+            download_zip_dataset(dataset_url, args.output_dir)
+            process_dataset(dataset_name, args.ou)

scripts/download_egfx.sh

@@ -1,22 +0,0 @@
-#/bin/bash
-mkdir -p data
-cd data
-mkdir -p egfx
-cd egfx
-wget https://zenodo.org/record/7044411/files/BluesDriver.zip?download=1 -O BluesDriver.zip
-wget https://zenodo.org/record/7044411/files/Chorus.zip?download=1 -O Chorus.zip
-wget https://zenodo.org/record/7044411/files/Clean.zip?download=1 -O Clean.zip
-wget https://zenodo.org/record/7044411/files/Digital-Delay.zip?download=1 -O Digital-Delay.zip
-wget https://zenodo.org/record/7044411/files/Flanger.zip?download=1 -O Flanger.zip
-wget https://zenodo.org/record/7044411/files/Hall-Reverb.zip?download=1 -O Hall-Reverb.zip
-wget https://zenodo.org/record/7044411/files/Phaser.zip?download=1 -O Phaser.zip
-wget https://zenodo.org/record/7044411/files/Plate-Reverb.zip?download=1 -O Plate-Reverb.zip
-wget https://zenodo.org/record/7044411/files/RAT.zip?download=1 -O RAT.zip
-wget https://zenodo.org/record/7044411/files/Spring-Reverb.zip?download=1 -O Spring-Reverb.zip
-wget https://zenodo.org/record/7044411/files/Sweep-Echo.zip?download=1 -O Sweep-Echo.zip
-wget https://zenodo.org/record/7044411/files/TapeEcho.zip?download=1 -O TapeEcho.zip
-wget https://zenodo.org/record/7044411/files/TubeScreamer.zip?download=1 -O TubeScreamer.zip
-unzip -n \*.zip
-rm -rf *.zip
-
-

scripts/test.py

@@ -2,7 +2,6 @@ import pytorch_lightning as pl
 import hydra
 from omegaconf import DictConfig
 import remfx.utils as utils
-from pytorch_lightning.utilities.model_summary import ModelSummary
 import torch
 
 log = utils.get_logger(__name__)