Remove unneeded scripts. Change eval to use table 4 datasets. Clean

README.md

@@ -4,29 +4,29 @@ Removing multiple audio effects from multiple sources using compositional audio
 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)
 
 
-
 # Setup
 ```
 git clone https://github.com/mhrice/RemFx.git
+cd RemFx
 git submodule update --init --recursive
 pip install -e . ./umx
 ```
 # Usage
 This repo can be used for many different tasks. Here are some examples.
-## Run RemFX Detect on a single file
+## Run RemFX Detect on a single file - []
+First, need to download the checkpoints from [zenodo](https://zenodo.org/record/8179396)
 ```
 ./download_checkpoints.sh
 ./remfx_detect.sh wet.wav -o dry.wav
 ```
-## Download the [General Purpose Audio Effect Removal evaluation dataset](https://zenodo.org/record/8183649/)
+## Download the [General Purpose Audio Effect Removal evaluation datasets](https://zenodo.org/record/8183649/) - [x]
 ```
-wget https://zenodo.org/record/8183649/files/RemFX_eval_dataset.zip?download=1 -O RemFX_eval_dataset.zip
-unzip RemFX_eval_dataset.zip
+./download_eval_datasets.sh
 ```
 
-## Download the starter datasets
+## Download the starter datasets - [x]
 ```
-python scripts/download.py vocalset guitarset idmt-smt-bass idmt-smt-drums
+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`
 
@@ -35,7 +35,7 @@ Then set the dataset root :
 export DATASET_ROOT={path/to/datasets}
 ```
 
-## Training
+## Training - [x]
 Before training, it is important that you have downloaded the starter 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
 ```
@@ -44,13 +44,13 @@ 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          |
-| ----------------------- | ------------ | ---------------- |
-| Effect-specific         | {effect}     | +exp=chorus      |
-| Effect-specific + FXAug | {effect}_aug | +exp=chorus_aug  |
-| Monolithic (1 FX)       | 5-5          | +exp=5-1         |
-| Monolithic (<=5 FX)     | 5-5          | +exp=5-5         |
-| Classifier              | 5-5_cls      | +exp=5-5_cls     |
+| Experiment Type         | Config Name  | Example           |
+| ----------------------- | ------------ | ----------------- |
+| Effect-specific         | {effect}     | +exp=chorus       |
+| Effect-specific + FXAug | {effect}_aug | +exp=chorus_aug   |
+| Monolithic (1 FX)       | 5-1          | +exp=5-1          |
+| 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.
 You can also create a custom experiment by creating a new experiment file in `cfg/exp/` and overriding the default parameters in `config.yaml`.
@@ -58,7 +58,7 @@ You can also create a custom experiment by creating a new experiment file in `cf
 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 test.py +exp={experiment_name} ckpt_path={path/to/checkpoint}
+python scripts/test.py +exp={experiment_name} +ckpt_path={path/to/checkpoint} render_files=False
 ```
 
 The checkpoints will be saved in `./logs/ckpts/{timestamp}`
@@ -69,27 +69,43 @@ If you have generated the dataset separately (see Generate datasets used in the
 
 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).
+## 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
 ```
 ./download_checkpoints.sh
 ```
-Then run the evaluation script, select the RemFX configuration, between `remfx_oracle`, `remfx_detect`, and `remfx_all`.
+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.
 ```
-./eval.sh remfx_detect
+./eval.sh remfx_detect 0-0
+./eval.sh remfx_detect 1-1
+./eval.sh remfx_detect 2-2
+./eval.sh remfx_detect 3-3
+./eval.sh remfx_detect 4-4
+./eval.sh remfx_detect 5-5
+
 ```
-To use a custom trained model, first train a model (see Training)
-Then run the evaluation script, with config used.
+To eval a custom monolithic model, first train a model (see Training)
+Then run the evaluation script, with the config used and checkpoint_path.
 ```
-./eval.sh {experiment_name}
+./eval.sh distortion_aug 0-0 -ckpt "logs/ckpts/2023-07-26-10-10-27/epoch\=05-valid_loss\=8.623.ckpt"
+./eval.sh distortion_aug 1-1 -ckpt "logs/ckpts/2023-07-26-10-10-27/epoch\=05-valid_loss\=8.623.ckpt"
+./eval.sh distortion_aug 2-2 -ckpt "logs/ckpts/2023-07-26-10-10-27/epoch\=05-valid_loss\=8.623.ckpt"
+./eval.sh distortion_aug 3-3 -ckpt "logs/ckpts/2023-07-26-10-10-27/epoch\=05-valid_loss\=8.623.ckpt"
+./eval.sh distortion_aug 4-4 -ckpt "logs/ckpts/2023-07-26-10-10-27/epoch\=05-valid_loss\=8.623.ckpt"
+./eval.sh distortion_aug 5-5 -ckpt "logs/ckpts/2023-07-26-10-10-27/epoch\=05-valid_loss\=8.623.ckpt"
 ```
 
-## Checkpoints
-Download checkpoints from [here](https://zenodo.org/record/8179396), or see the ./download_checkpoints.sh script.
+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}.
+Then run the evaluation script.
+```
+./eval.sh remfx_detect 0-0
+```
 
+The script assumes that RemFX_eval_datasets is in the top-level directory.
+Metrics and hyperparams will be logged in `./lightning_logs/{timestamp}`
 
-## Generate datasets used in the paper
+## Generate other datasets - [x]
 The datasets used in the experiments are customly generated from the starter datasets. In short, for each training/val/testing example, we select a random 5.5s segment from one of the starter datasets and apply a random number of effects to it. The number of effects applied is controlled by the `num_kept_effects` and `num_removed_effects` parameters. The effects applied are controlled by the `effects_to_keep` and `effects_to_remove` parameters.
 
 Before generating datasets, it is important that you have downloaded the starter datasets (see above) and set DATASET_ROOT.
@@ -105,28 +121,6 @@ By default, files are rendered to `{render_root} / processed / {string_of_effect
 
 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.
 
-## Evaluate with a custom directory
-Assumes directory is structured as
-- root
-    - clean
-        - file1.wav
-        - file2.wav
-        - file3.wav
-    - effected
-        - file1.wav
-        - file2.wav
-        - file3.wav
-
-First set the dataset root:
-```
-export DATASET_ROOT={path/to/datasets}
-```
-
-Then run
-```
-python scripts/chain_inference.py +exp=chain_inference_custom
-```
-
 # Misc.
 ## Experimental parameters
 Some relevant dataset/training parameters descriptions
@@ -159,3 +153,26 @@ Some relevant dataset/training parameters descriptions
 - `distortion`
 - `reverb`
 - `delay`
+
+# DO WE NEED THIS?
+## Evaluate RemFXwith a custom directory - []
+Assumes directory is structured as
+- root
+    - clean
+        - file1.wav
+        - file2.wav
+        - file3.wav
+    - effected
+        - file1.wav
+        - file2.wav
+        - file3.wav
+
+First set the dataset root:
+```
+export DATASET_ROOT={path/to/datasets}
+```
+
+Then run
+```
+python scripts/chain_inference.py +exp=chain_inference_custom
+```

cfg/exp/0-0.yaml

@@ -0,0 +1,29 @@
+# @package _global_
+defaults:
+  - override /model: demucs
+  - override /effects: all
+seed: 12345
+sample_rate: 48000
+chunk_size: 262144 # 5.5s
+logs_dir: "./logs"
+render_files: True
+
+accelerator: "gpu"
+log_audio: True
+# Effects
+num_kept_effects: [0,0] # [min, max]
+num_removed_effects: [0,0] # [min, max]
+shuffle_kept_effects: True
+shuffle_removed_effects: True
+num_classes: 5
+effects_to_keep:
+effects_to_remove:
+  - distortion
+  - compressor
+  - reverb
+  - chorus
+  - delay
+datamodule:
+  train_batch_size: 16
+  test_batch_size: 1
+  num_workers: 8

cfg/exp/1-1.yaml

@@ -12,7 +12,7 @@ accelerator: "gpu"
 log_audio: True
 # Effects
 num_kept_effects: [0,0] # [min, max]
-num_removed_effects: [0,1] # [min, max]
+num_removed_effects: [1,1] # [min, max]
 shuffle_kept_effects: True
 shuffle_removed_effects: True
 num_classes: 5

cfg/exp/2-2.yaml

@@ -12,7 +12,7 @@ accelerator: "gpu"
 log_audio: True
 # Effects
 num_kept_effects: [0,0] # [min, max]
-num_removed_effects: [0,2] # [min, max]
+num_removed_effects: [2,2] # [min, max]
 shuffle_kept_effects: True
 shuffle_removed_effects: True
 num_classes: 5

cfg/exp/3-3.yaml

@@ -12,7 +12,7 @@ accelerator: "gpu"
 log_audio: True
 # Effects
 num_kept_effects: [0,0] # [min, max]
-num_removed_effects: [0,3] # [min, max]
+num_removed_effects: [3,3] # [min, max]
 shuffle_kept_effects: True
 shuffle_removed_effects: True
 num_classes: 5

cfg/exp/4-4.yaml

@@ -12,7 +12,7 @@ accelerator: "gpu"
 log_audio: True
 # Effects
 num_kept_effects: [0,0] # [min, max]
-num_removed_effects: [0,4] # [min, max]
+num_removed_effects: [4,4] # [min, max]
 shuffle_kept_effects: True
 shuffle_removed_effects: True
 num_classes: 5

cfg/exp/5-5.yaml

@@ -12,7 +12,7 @@ accelerator: "gpu"
 log_audio: True
 # Effects
 num_kept_effects: [0,0] # [min, max]
-num_removed_effects: [0,5] # [min, max]
+num_removed_effects: [5,5] # [min, max]
 shuffle_kept_effects: True
 shuffle_removed_effects: True
 num_classes: 5

cfg/exp/5-5_full.yaml

@@ -0,0 +1,29 @@
+# @package _global_
+defaults:
+  - override /model: demucs
+  - override /effects: all
+seed: 12345
+sample_rate: 48000
+chunk_size: 262144 # 5.5s
+logs_dir: "./logs"
+render_files: True
+
+accelerator: "gpu"
+log_audio: True
+# Effects
+num_kept_effects: [0,0] # [min, max]
+num_removed_effects: [0,5] # [min, max]
+shuffle_kept_effects: True
+shuffle_removed_effects: True
+num_classes: 5
+effects_to_keep:
+effects_to_remove:
+  - distortion
+  - compressor
+  - reverb
+  - chorus
+  - delay
+datamodule:
+  train_batch_size: 16
+  test_batch_size: 1
+  num_workers: 8

cfg/exp/remfx_all.yaml

@@ -6,7 +6,7 @@ seed: 12345
 sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
-accelerator: "cpu"
+accelerator: "gpu"
 log_audio: True
 
 # Effects
@@ -17,11 +17,11 @@ shuffle_removed_effects: True
 num_classes: 5
 effects_to_keep:
 effects_to_remove:
+  - distortion
   - compressor
   - reverb
   - chorus
   - delay
-  - distortion
 datamodule:
   train_batch_size: 16
   test_batch_size: 1
@@ -85,6 +85,4 @@ inference_effects_ordering:
   - "RandomPedalboardDelay"
 num_bins: 1025
 inference_effects_shuffle: True
-inference_use_all_effect_models: True
-audio_input: ""
-output_path: "./output.wav"
+inference_use_all_effect_models: True

cfg/exp/remfx_detect.yaml

@@ -6,7 +6,7 @@ seed: 12345
 sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
-accelerator: "cpu"
+accelerator: "gpu"
 log_audio: True
 
 # Effects
@@ -17,11 +17,11 @@ shuffle_removed_effects: True
 num_classes: 5
 effects_to_keep:
 effects_to_remove:
+  - distortion
   - compressor
   - reverb
   - chorus
   - delay
-  - distortion
 datamodule:
   train_batch_size: 16
   test_batch_size: 1
@@ -85,6 +85,4 @@ inference_effects_ordering:
   - "RandomPedalboardDelay"
 num_bins: 1025
 inference_effects_shuffle: True
-inference_use_all_effect_models: False
-audio_input: ""
-output_path: "./output.wav"
+inference_use_all_effect_models: False

cfg/exp/remfx_oracle.yaml

@@ -6,7 +6,7 @@ seed: 12345
 sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
-accelerator: "cpu"
+accelerator: "gpu"
 log_audio: True
 
 # Effects
@@ -17,11 +17,11 @@ shuffle_removed_effects: True
 num_classes: 5
 effects_to_keep:
 effects_to_remove:
+  - distortion
   - compressor
   - reverb
   - chorus
   - delay
-  - distortion
 datamodule:
   train_batch_size: 16
   test_batch_size: 1
@@ -69,6 +69,4 @@ inference_effects_ordering:
   - "RandomPedalboardDelay"
 num_bins: 1025
 inference_effects_shuffle: True
-inference_use_all_effect_models: False
-audio_input: ""
-output_path: "./output.wav"
+inference_use_all_effect_models: False

download_ckpts.sh

@@ -1,3 +1,5 @@
+#! /bin/bash
+
 # make ckpts directory if not exist
 mkdir -p ckpts
 

download_eval_datasets.sh

@@ -0,0 +1,25 @@
+#! /bin/bash
+
+mkdir -p RemFX_eval_datasets
+cd RemFX_eval_datasets
+mkdir -p processed
+cd processed
+wget https://zenodo.org/record/8187288/files/0-0.zip?download=1 -O 0-0.zip
+wget https://zenodo.org/record/8187288/files/1-1.zip?download=1 -O 1-1.zip
+wget https://zenodo.org/record/8187288/files/2-2.zip?download=1 -O 2-2.zip
+wget https://zenodo.org/record/8187288/files/3-3.zip?download=1 -O 3-3.zip
+wget https://zenodo.org/record/8187288/files/4-4.zip?download=1 -O 4-4.zip
+wget https://zenodo.org/record/8187288/files/5-5.zip?download=1 -O 5-5.zip
+unzip 0-0.zip
+unzip 1-1.zip
+unzip 2-2.zip
+unzip 3-3.zip
+unzip 4-4.zip
+unzip 5-5.zip
+rm 0-0.zip
+rm 1-1.zip
+rm 2-2.zip
+rm 3-3.zip
+rm 4-4.zip
+rm 5-5.zip
+

eval.sh

@@ -1,16 +1,51 @@
 #! /bin/bash
 
 # Example usage:
-# ./eval.sh remfx_detect 
+# ./eval.sh remfx_detect 0-0
+# ./eval.sh distortion_aug 0-0 -ckpt logs/ckpts/2023-01-21-12-21-44
+# First 2 arguments are required, third argument is optional
 
 # Check if first argument is empty
 if [ -z "$1" ]
 then
-  echo "No experiment name or config path supplied"
+  echo "No experiment name supplied"
   exit 1
 fi
 
-python scripts/chain_inference.py +exp=$1 datamodule.train_dataset=None datamodule.val_dataset=None datamodule.test_dataset.render_root=./RemFX_eval_dataset/ render_files=False
+# Check if second argument is empty
+if [ -z "$2" ]
+then
+  echo "No dataset name supplied"
+  exit 1
+fi
+dataset_name=$2
+
+# Check if ckpt flag is set using getopts
+ckpt_flag=0
+while getopts ":ckpt:" opt; do
+  case $opt in
+    ckpt)
+      ckpt_flag=1
+      ckpt_path=$OPTARG
+      ;;
+    \?)
+      echo "Invalid option: -$OPTARG" >&3
+      ;;
+  esac
+done
+
+# If checkpoint flag is empty, run chain inference
+if [ $ckpt_flag -eq 0 ]
+then
+  # Running chain inference
+  echo "Running chain inference"
+  python scripts/chain_inference.py +exp=$1 datamodule.train_dataset=None datamodule.val_dataset=None datamodule.test_dataset.render_root=./RemFX_eval_datasets/ render_files=False num_removed_effects=[${dataset_name:0:1},${dataset_name:2:1}]
+  exit 1
+fi
+
+# Otherwise run inference on the specified checkpoint
+echo "Running monolithic inference on checkpoint $3"
+python scripts/test.py +exp=$1 datamodule.train_dataset=None datamodule.val_dataset=None datamodule.test_dataset.render_root=./RemFX_eval_datasets/ datamodule.test_dataset.num_kept_effects="[0,0]" num_removed_effects=[${dataset_name:0:1},${dataset_name:2:1}] effects_to_keep=[] effects_to_remove="[compressor,reverb,chorus,delay,distortion]" render_files=False +ckpt_path=$2
 
 
 

remfx/datasets.py

@@ -578,6 +578,9 @@ class EffectDataset(Dataset):
             normalized_wet = self.normalize(wet)
 
             # Check STFT, pick different effects if necessary
+            if num_removed_effects == 0:
+                # No need to check if no effects removed
+                break
             stft = self.mrstft(normalized_wet.unsqueeze(0), normalized_dry.unsqueeze(0))
         return normalized_dry, normalized_wet, dry_labels_tensor, wet_labels_tensor
 

remfx/utils.py

@@ -3,7 +3,6 @@ from typing import List, Tuple
 import pytorch_lightning as pl
 from omegaconf import DictConfig
 from pytorch_lightning.utilities import rank_zero_only
-import numpy as np
 import torch
 import torchaudio
 from torch import nn

remfx_detect.sh

@@ -33,7 +33,7 @@ done
 
 if [ -z "$output_path" ]
 then
-  python scripts/remfx_detect.py +exp=remfx_detect audio_input=$1
+  python scripts/remfx_detect.py +exp=remfx_detect +audio_input=$audio_input
   exit 0
 fi
-python scripts/remfx_detect.py +exp=remfx_detect audio_input=$1 output_path=$output_path
+python scripts/remfx_detect.py +exp=remfx_detect +audio_input=$audio_input +output_path=$output_path

scripts/remfx_detect.py

@@ -39,7 +39,7 @@ def main(cfg: DictConfig):
         use_all_effect_models=cfg.inference_use_all_effect_models,
     )
 
-    audio_file = "/Users/matthewrice/Desktop/clips/chipmunk.wav"
+    audio_file = cfg.audio_input
     print("Loading", audio_file)
     audio, sr = torchaudio.load(audio_file)
     # Resample
@@ -51,8 +51,12 @@ def main(cfg: DictConfig):
     batch = [audio, audio, None, None]
 
     _, y = inference_model(batch, 0, verbose=True)
-    print("Saving output to", cfg.output_path)
-    torchaudio.save(cfg.output_path, y[0], sample_rate=cfg.sample_rate)
+    if "output_path" in cfg:
+        output_path = cfg.output_path
+    else:
+        output_path = "./output.wav"
+    print("Saving output to", output_path)
+    torchaudio.save(output_path, y[0], sample_rate=cfg.sample_rate)
 
 
 if __name__ == "__main__":

shell_vars.sh

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

train_all.sh

@@ -1,6 +0,0 @@
-python scripts/train.py +exp=5-5_cls.yaml model=cls_wav2vec2 render_files=False logs_dir=/scratch/cjs-log
-python scripts/train.py +exp=5-5_cls.yaml model=cls_panns_44k render_files=False logs_dir=/scratch/cjs-log
-python scripts/train.py +exp=5-5_cls.yaml model=cls_panns_16k render_files=False logs_dir=/scratch/cjs-log 
-python scripts/train.py +exp=5-5_cls.yaml model=cls_panns_pt render_files=False logs_dir=/scratch/cjs-log
-python scripts/train.py +exp=5-5_cls.yaml model=cls_vggish render_files=False logs_dir=/scratch/cjs-log 
-python scripts/train.py +exp=5-5_cls.yaml model=cls_wav2clip render_files=False logs_dir=/scratch/cjs-log