Update README, fix minor instance checks

README.md

@@ -1,7 +1,20 @@
 # General Purpose Audio Effect Removal
-Removing multiple audio effects from multiple sources using compositional audio effect removal and source separation and speech enhancement models.
+Removing multiple audio effects from multiple sources with compositional audio effect removal using 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, citation, license)
+This repo contains the code for the paper [General Purpose Audio Effect Removal](https://arxiv.org/abs/2110.00484). (Todo: Paper link broken, Arxiv badge broken, citation, license)
+
+[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1LoLgL1YHzIQfILEayDmRUZzDZzJpD6rD)
+[![arXiv](https://img.shields.io/badge/arXiv-1234.56789-b31b1b.svg)](https://arxiv.org/abs/1234.56789)
+
+![RemFX Block diagram](remfx-headline.jpg?raw=true "Title")
+
+Listening examples can be found [here](https://csteinmetz1.github.io/RemFX/).
+
+## Abstract
+
+Although the design and application of audio effects is well understood, the inverse problem of removing these effects is significantly more challenging and far less studied. Recently, deep learning has been applied to audio effect removal; however, existing approaches have focused on narrow formulations considering only one effect or source type at a time. In realistic scenarios, multiple effects are applied with varying source content. This motivates a more general task, which we refer to as general purpose audio effect removal. We developed a dataset for this task using five audio effects across four different sources and used it to train and evaluate a set of existing architectures. We found that no single model performed optimally on all effect types and sources. To address this, we introduced <b>RemFX</b>, an approach designed to mirror the compositionality of applied effects. We first trained a set of the best-performing effect-specific
+removal models and then leveraged an audio effect classification model to dynamically construct a graph of our models at inference. We found our approach to outperform single model baselines, although examples with many effects present remain challenging.
 
 
 # Setup
@@ -13,29 +26,38 @@ pip install -e . ./umx
 pip install --no-deps hearbaseline
 ```
 Due to incompatabilities with hearbaseline's dependencies (namely numpy/numba) and our other packages, we need to install hearbaseline with no dependencies.
+<b>Please run the setup code before running any scripts.</b>
+All scripts should be launched from the top level after installing.
+
 # Usage
-This repo can be used for many different tasks. Here are some examples.
+This repo can be used for many different tasks. Here are some examples. Ensure you have run the setup code before running any scripts.
 ## Run RemFX Detect on a single file
-First, need to download the checkpoints from [zenodo](https://zenodo.org/record/8179396)
+Here we will attempt to detect, then remove effects that are present in an audio file. For the best results, use a file from our [evaluation dataset](https://zenodo.org/record/8187288). We support detection and removal of the following effects: chorus, delay, distortion, dynamic range compression, and reverb.
+
+First, we need to download the pytorch checkpoints from [zenodo](https://zenodo.org/record/8218621)
 ```
-scripts/download_checkpoints.sh
+scripts/download_ckpts.sh
 ```
 Then run the detect script. This repo contains an example file `example.wav` from our test dataset which contains 2 effects (chorus and delay) applied to a guitar.
 ```
 scripts/remfx_detect.sh example.wav -o dry.wav
 ```
 ## Download the [General Purpose Audio Effect Removal evaluation datasets](https://zenodo.org/record/8187288)
+We provide a script to download and unzip the datasets used in table 4 of the paper.
 ```
 scripts/download_eval_datasets.sh
 ```
 
 ## Download the starter datasets
+
+If you'd like to train your own model and/or generate a dataset, you can download the starter datasets using the following command:
+
 ```
 python scripts/download.py vocalset guitarset dsd100 idmt-smt-drums
 ```
 By default, the starter datasets are downloaded to `./data/remfx-data`. To change this, pass `--output_dir={path/to/datasets}` to `download.py`
 
-Then set the dataset root :
+Then set the dataset root:
 ```
 export DATASET_ROOT={path/to/datasets}
 ```
@@ -47,6 +69,13 @@ This project uses the [pytorch-lightning](https://www.pytorchlightning.ai/index.
 python scripts/train.py +exp={experiment_name}
 ```
 
+At the end of training, the train script will automatically evaluate the test set using the best checkpoint (by validation loss). If epoch 0 is not finished, it will throw an error. To evaluate a specific checkpoint, run
+
+```
+python scripts/test.py +exp={experiment_name} +ckpt_path="{path/to/checkpoint}" render_files=False
+```
+
+### Experiments
 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           |
@@ -57,42 +86,39 @@ Here are some selected experiment types from the paper, which use different data
 | Monolithic (<=5 FX)     | 5-5_full     | +exp=5-5_full     |
 | Classifier              | 5-5_full_cls | +exp=5-5_full_cls |
 
-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.
+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 Experimental parameters 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). If epoch 0 is not finished, it will throw an error. To evaluate a specific checkpoint, run
+### Logging
+By default, training uses the Pytorch Lightning CSV Logger
+Metrics and hyperparams will be logged in `./lightning_logs/{timestamp}`
+
+[Weights and Biases](https://wandb.ai/) logging can also be used, and will log audio during training and testing. To use Weights and Biases, set `logger=wandb` in the config or command-line. Make sure you have an account and are logged in.
 
+Then set the project and entity:
 ```
-python scripts/test.py +exp={experiment_name} +ckpt_path="{path/to/checkpoint}" render_files=False
+export WANDB_PROJECT={desired_wandb_project}
+export WANDB_ENTITY={your_wandb_username}
 ```
 
 The checkpoints will be saved in `./logs/ckpts/{timestamp}`
-Metrics and hyperparams will be logged in `./lightning_logs/{timestamp}`
 
-By default, the dataset needed for the experiment is generated before training.
+### Misc.
+- By default, the dataset needed for the experiment is generated before training.
 If you have generated the dataset separately (see Generate datasets used in the paper), 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.
-
-### Logging
-Default CSV logger
-To use WANDB logger:
+- Training assumes you have a CUDA GPU. To train on CPU, set `accelerator=null` in the config or command-line.
 
-export WANDB_PROJECT={desired_wandb_project}
-export WANDB_ENTITY={your_wandb_username}
+- If training with the pretrained PANNs model, download the pretrained model from [here](https://zenodo.org/record/6332525) or run: `wget https://zenodo.org/record/6332525/files/hear2021-panns_hear.pth`. Place this in the root of the repo.
 
-## Panns pretrianed
-```
-wget https://zenodo.org/record/6332525/files/hear2021-panns_hear.pth
-```
 
 ## Evaluate models on the General Purpose Audio Effect Removal evaluation datasets (Table 4 from the paper)
-First download the General Purpose Audio Effect Removal evaluation datasets (see above).
-To use the pretrained RemFX model, download the checkpoints
+We provide a way to replicate the results of table 4 from our paper. First download the <b>General Purpose Audio Effect Removal evaluation datasets</b> (see above).
+To use the pretrained RemFX model, download the checkpoints:
 ```
-scripts/download_checkpoints.sh
+scripts/download_ckpts.sh
 ```
-Then run the evaluation script, select the RemFX configuration, between `remfx_oracle`, `remfx_detect`, and `remfx_all`. Then select N, the number of effects to remove.
+Then run the evaluation script. First select the RemFX configuration, between `remfx_oracle`, `remfx_detect`, and `remfx_all`. As a reminder, `remfx_oracle` uses the ground truth labels of the present effects to determine which removal models to apply, `remfx_detect` detects which effects are present, and `remfx_all` assumes all effects are present.
 ```
 scripts/eval.sh remfx_detect 0-0
 scripts/eval.sh remfx_detect 1-1
@@ -100,15 +126,17 @@ scripts/eval.sh remfx_detect 2-2
 scripts/eval.sh remfx_detect 3-3
 scripts/eval.sh remfx_detect 4-4
 scripts/eval.sh remfx_detect 5-5
-
 ```
+In this case the `N-N` refers to the number of effects present for each example in the dataset.
+
+
 To eval a custom monolithic model, first train a model (see Training)
 Then run the evaluation script, with the config used and checkpoint_path.
 ```
-scripts/eval.sh distortion_aug 0-0 -ckpt "logs/ckpts/2023-07-26-10-10-27/epoch\=05-valid_loss\=8.623.ckpt"
+scripts/eval.sh distortion_aug 0-0 -ckpt "{path/to/checkpoint}"
 ```
 
-To eval a custom effect-specific model as part of the inference chain, first train a model (see Training), then edit `cfg/exp/remfx_{desired_configuration}.yaml -> ckpts -> {effect}`.
+To eval a custom effect-specific model as part of the inference chain, first train a model (see Training), then edit `cfg/exp/remfx_{desired_configuration}.yaml -> ckpts -> {effect}`. Select between `remfx_detect`, `remfx_oracle`, and `remfx_all`.
 Then run the evaluation script.
 ```
 scripts/eval.sh remfx_detect 0-0
@@ -128,40 +156,61 @@ For example, to generate the `chorus` FXAug dataset, which includes files with 5
 python scripts/generate_dataset.py +exp=chorus_aug
 ```
 
-See the Misc. section below for a description of the parameters.
+See the Experimental parameters section below for a description of the parameters.
 By default, files are rendered to `{render_root} / processed / {string_of_effects} / {train|val|test}`.
 
-If training, this process will be done automatically at the start of training. To disable this, set `render_files=False` in the config or command-line, and set `render_root={path/to/dataset}` if it is in a custom location.
-
-# Misc.
-## Experimental parameters
+The dataset that is generated contains 8000 train examples, 1000 validation examples, and 1000 test examples. Each example is contained in a folder labeled by its id number (ex. 0-7999 for train examples) with 4 files like so:
+```
+.
+└── train
+    ├── 0
+    │   ├── dry_effects.pt
+    │   ├── input.wav
+    │   ├── target.wav
+    │   └── wet_effects.pt
+    ├── 1
+    │   └── ...
+    ├── ...
+    ├── 7999
+    │   └── ...
+```
+### File descriptions
+- dry_effects.pt = serialized PyTorch file that contains a list of the effects applied to the dry audio file
+- input.wav = the wet audio file
+- target.wav = the dry audio file
+- wet_effects.pt = serialized PyTorch file that contains a list of the effects applied to the wet audio file
+
+Note: if training, this process will be done automatically at the start of training. To disable this, set `render_files=False` in the config or command-line, and set `render_root={path/to/dataset}` if it is in a custom location.
+
+# Experimental parameters
 Some relevant dataset/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 'Effect Removal Models/Effect Classification 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'). Used for FXAug.
-- `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: ./data)
-- `datamodule.train_batch_size={batch_size}`. Change batch size (default: varies)
-
-### Effect Removal Models
+- `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: ./data).
+- `datamodule.train_batch_size={batch_size}`. Change batch size (default: varies).
+- `logger=wandb`. Use weights and biases logger (default: csv). Ensure you set the wandb environment variables (see training section).
+
+## Effect Removal Models
 - `umx`
 - `demucs`
 - `tcn`
 - `dcunet`
 - `dptnet`
 
-### Effect Classification Models
+## Effect Classification Models
 - `cls_vggish`
 - `cls_panns_pt`
 - `cls_wav2vec2`
 - `cls_wav2clip`
 
-### Effects
+## Effects
+- `delay`
+- `distortion`
 - `chorus`
 - `compressor`
-- `distortion`
 - `reverb`
-- `delay`

remfx/callbacks.py

@@ -90,7 +90,7 @@ def log_wandb_audio_batch(
     caption: str = "",
     max_items: int = 10,
 ):
-    if type(logger) != pl.loggers.WandbLogger:
+    if not isinstance(logger, pl.loggers.WandbLogger):
         return
     num_items = samples.shape[0]
     samples = rearrange(samples, "b c t -> b t c")

remfx/utils.py

@@ -72,7 +72,7 @@ def log_hyperparameters(
     if "callbacks" in config:
         hparams["callbacks"] = config["callbacks"]
 
-    if type(trainer.logger) == pl.loggers.CSVLogger:
+    if isinstance(logger, pl.loggers.CSVLogger):
         logger.log_hyperparams(hparams)
     else:
         logger.experiment.config.update(hparams)

scripts/download_ckpts.sh

@@ -4,9 +4,9 @@
 mkdir -p ckpts
 
 # download ckpts and save to ckpts directory
-wget https://zenodo.org/record/8179396/files/classifier.ckpt?download=1 -O ckpts/classifier.ckpt
-wget https://zenodo.org/record/8179396/files/dcunet_chorus_aug.ckpt?download=1 -O ckpts/dcunet_chorus_aug.ckpt
-wget https://zenodo.org/record/8179396/files/dcunet_delay_aug.ckpt?download=1 -O ckpts/dcunet_delay_aug.ckpt
-wget https://zenodo.org/record/8179396/files/dcunet_reverb_aug.ckpt?download=1 -O ckpts/dcunet_reverb_aug.ckpt
-wget https://zenodo.org/record/8179396/files/demucs_compressor_aug.ckpt?download=1 -O ckpts/demucs_compressor_aug.ckpt
-wget https://zenodo.org/record/8179396/files/demucs_distortion_aug.ckpt?download=1 -O ckpts/demucs_distortion_aug.ckpt
+wget https://zenodo.org/record/8218621/files/classifier.ckpt?download=1 -O ckpts/classifier.ckpt
+wget https://zenodo.org/record/8218621/files/dcunet_chorus_aug.ckpt?download=1 -O ckpts/dcunet_chorus_aug.ckpt
+wget https://zenodo.org/record/8218621/files/dcunet_delay_aug.ckpt?download=1 -O ckpts/dcunet_delay_aug.ckpt
+wget https://zenodo.org/record/8218621/files/dcunet_reverb_aug.ckpt?download=1 -O ckpts/dcunet_reverb_aug.ckpt
+wget https://zenodo.org/record/8218621/files/demucs_compressor_aug.ckpt?download=1 -O ckpts/demucs_compressor_aug.ckpt
+wget https://zenodo.org/record/8218621/files/demucs_distortion_aug.ckpt?download=1 -O ckpts/demucs_distortion_aug.ckpt

scripts/generate_dataset.py

@@ -8,7 +8,7 @@ def main(cfg: DictConfig):
     # Apply seed for reproducibility
     if cfg.seed:
         pl.seed_everything(cfg.seed)
-    datamodule = hydra.utils.instantiate(cfg.datamodule, _convert_="partial")
+    _ = hydra.utils.instantiate(cfg.datamodule, _convert_="partial")
 
 
 if __name__ == "__main__":