Addon

AudioSep_Colab.ipynb

@@ -0,0 +1,128 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "from pathlib import Path\n",
+        "\n",
+        "repo_path = Path(\"/content/AudioSep\")\n",
+        "if not repo_path.exists():\n",
+        "    !git clone https://github.com/Audio-AGI/AudioSep.git\n",
+        "\n",
+        "%cd /content/AudioSep"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "pjIhw5ECS_3_"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install torchlibrosa==0.1.0 gradio==3.47.1 gdown lightning transformers==4.28.1 ftfy braceexpand webdataset soundfile wget h5py"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "t6h9KB3CcjBd"
+      },
+      "outputs": [],
+      "source": [
+        "checkpoints_dir = Path(\"checkpoint\")\n",
+        "checkpoints_dir.mkdir(exist_ok=True)\n",
+        "\n",
+        "models = (\n",
+        "    (\n",
+        "        \"https://huggingface.co/spaces/badayvedat/AudioSep/resolve/main/checkpoint/audiosep_base_4M_steps.ckpt\",\n",
+        "        checkpoints_dir / \"audiosep_base_4M_steps.ckpt\"\n",
+        "    ),\n",
+        "    (\n",
+        "        \"https://huggingface.co/spaces/badayvedat/AudioSep/resolve/main/checkpoint/music_speech_audioset_epoch_15_esc_89.98.pt\",\n",
+        "        checkpoints_dir / \"music_speech_audioset_epoch_15_esc_89.98.pt\"\n",
+        "    )\n",
+        ")\n",
+        "\n",
+        "for model_url, model_path in models:\n",
+        "    if not model_path.exists():\n",
+        "        !wget {model_url} -O {model_path}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3uDrzCQyY58h"
+      },
+      "outputs": [],
+      "source": [
+        "!wget \"https://audio-agi.github.io/Separate-Anything-You-Describe/demos/exp31_water drops_mixture.wav\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0nr77CGXTwO1"
+      },
+      "outputs": [],
+      "source": [
+        "import torch\n",
+        "from pipeline import build_audiosep, inference\n",
+        "\n",
+        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
+        "\n",
+        "model = build_audiosep(\n",
+        "      config_yaml='config/audiosep_base.yaml',\n",
+        "      checkpoint_path=str(models[0][1]),\n",
+        "      device=device)\n",
+        "\n",
+        "audio_file = 'exp31_water drops_mixture.wav'\n",
+        "text = 'water drops'\n",
+        "output_file='separated_audio.wav'\n",
+        "\n",
+        "# AudioSep processes the audio at 32 kHz sampling rate\n",
+        "inference(model, audio_file, text, output_file, device)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "kssOe0pbPSWp"
+      },
+      "outputs": [],
+      "source": [
+        "print(f\"The separated audio is saved to: '{output_file}' file.\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "sl35U3dAR6KN"
+      },
+      "outputs": [],
+      "source": []
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

CONTRIBUTING.md

@@ -0,0 +1,92 @@
+# 🎵 Contributing to AudioSep
+
+Welcome to the AudioSep repository, where your contributions can harmonize the world of audio separation. To ensure a harmonious and organized collaboration, please follow the contribution guidelines outlined below.
+
+## **Submitting Contributions**
+
+To contribute to this project, please adhere to the following steps:
+
+### **1. Choose or Create an Issue**
+
+- Start by reviewing the existing issues to identify areas where your contributions can make a significant impact.
+- If you have ideas for new features, enhancements, or bug fixes, feel free to create a new issue to propose your contributions. Provide comprehensive details for clarity.
+
+### **2. Fork the Repository**
+
+- To initiate your contribution, fork the primary repository by clicking the "Fork" button. This will create a copy of the repository in your personal GitHub account.
+
+### **3. Clone Your Forked Repository**
+
+- Clone your forked repository to your local development environment using the following command:
+
+```bash
+git clone https://github.com/your-username/AudioSep.git
+```
+
+### **4. Set Up the Upstream Remote**
+
+- Maintain a reference to the primary project by adding it as the upstream remote:
+
+```bash
+cd AudioSep
+git remote add upstream https://github.com/Audio-AGI/AudioSep
+git remote -v
+```
+
+### **5. Create a New Branch**
+
+- Before starting your contribution, establish a new branch dedicated to your specific task:
+
+```bash
+git checkout -b my-contribution
+```
+
+## **Working on Your Contribution**
+
+Now that your development environment is ready and a new branch is established, you can start working on your contribution. Please ensure you adhere to the following guidelines:
+
+### **6. Make Changes**
+
+- Implement the necessary changes, including code additions, enhancements, or bug fixes. Ensure your contributions are well-structured, documented, and aligned with the project's objectives.
+
+### **7. Commit Your Changes**
+
+- Commit your changes using informative commit messages that clearly convey the purpose of your contributions:
+
+```bash
+git commit -m "Add a descriptive message here"
+```
+
+### **8. Push Your Changes**
+
+- Push the committed changes to your remote repository on GitHub:
+
+```bash
+git push origin my-contribution
+```
+
+### **9. Create a Pull Request**
+
+- Visit your repository on GitHub and click the "New Pull Request" button to initiate a pull request from your branch to the primary repository.
+
+### **10. Await Review**
+
+- Your pull request will undergo review, and feedback will be provided by the project maintainers or fellow contributors. Be prepared to address any suggested changes or refinements.
+
+## **Community Engagement**
+
+While contributing, please consider engaging with the community in the following ways:
+
+### **11. Join Discussions**
+
+- Participate in discussions related to audio separation techniques and their applications. Share your insights, experiences, and expertise in the audio field.
+
+### **12. Share Ideas**
+
+- If you have innovative ideas for advancing the project or optimizing audio separation, such as new algorithms or research findings, feel free to open issues to initiate productive discussions.
+
+## **Acknowledgment**
+
+We appreciate your dedication to the world of audio separation. Your contributions play a crucial role in harmonizing audio and improving the listening experience for all. If you have questions or require assistance, please don't hesitate to contact the project maintainers.
+
+Thank you for your valuable contributions, and we eagerly anticipate collaborating with you on AudioSep! 🎶🙌

Dockerfile

@@ -0,0 +1,22 @@
+FROM python:3.10.11
+
+# Copy the current directory contents into the container at .
+COPY . .
+
+# Set the working directory to /
+WORKDIR /
+
+# Install requirements.txt 
+RUN pip install --no-cache-dir --upgrade -r /requirements.txt
+
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH
+
+WORKDIR $HOME/app
+
+COPY --chown=user . $HOME/app
+
+# Start the FastAPI app on port 7860, the default port expected by Spaces
+CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "7860"]

LICENSE

@@ -0,0 +1,21 @@
+    MIT License
+
+    Copyright (c) Xubo Liu
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy
+    of this software and associated documentation files (the "Software"), to deal
+    in the Software without restriction, including without limitation the rights
+    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the Software is
+    furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all
+    copies or substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE

benchmark.py

@@ -0,0 +1,116 @@
+import os
+from tqdm import tqdm
+import numpy as np
+from evaluation.evaluate_audioset import AudioSetEvaluator
+from evaluation.evaluate_audiocaps import AudioCapsEvaluator
+from evaluation.evaluate_vggsound import VGGSoundEvaluator
+from evaluation.evaluate_music import MUSICEvaluator
+from evaluation.evaluate_esc50 import ESC50Evaluator
+from evaluation.evaluate_clotho import ClothoEvaluator
+from models.clap_encoder import CLAP_Encoder
+
+from utils import (
+    load_ss_model,
+    calculate_sdr,
+    calculate_sisdr,
+    parse_yaml,
+    get_mean_sdr_from_dict,
+)
+
+def eval(checkpoint_path, config_yaml='config/audiosep_base.yaml'):
+
+    log_dir = 'eval_logs'
+    os.makedirs(log_dir, exist_ok=True)
+
+    device = "cuda"
+    
+    configs = parse_yaml(config_yaml)
+
+    # AudioSet Evaluators
+    audioset_evaluator = AudioSetEvaluator()
+    # AudioCaps Evaluator
+    audiocaps_evaluator = AudioCapsEvaluator()
+    # VGGSound+ Evaluator
+    vggsound_evaluator = VGGSoundEvaluator()
+    # Clotho Evaluator
+    clotho_evaluator = ClothoEvaluator()
+    # MUSIC Evaluator
+    music_evaluator = MUSICEvaluator()
+    # ESC-50 Evaluator
+    esc50_evaluator = ESC50Evaluator()
+    
+    # Load model
+    query_encoder = CLAP_Encoder().eval()
+
+    pl_model = load_ss_model(
+        configs=configs,
+        checkpoint_path=checkpoint_path,
+        query_encoder=query_encoder
+    ).to(device)
+
+    print(f'-------  Start Evaluation  -------')
+
+    # evaluation on Clotho
+    SISDR, SDRi = clotho_evaluator(pl_model)
+    msg_clotho = "Clotho Avg SDRi: {:.3f}, SISDR: {:.3f}".format(SDRi, SISDR)
+    print(msg_clotho)
+    
+    # evaluation on VGGSound+ (YAN)
+    SISDR, SDRi = vggsound_evaluator(pl_model)
+    msg_vgg = "VGGSound Avg SDRi: {:.3f}, SISDR: {:.3f}".format(SDRi, SISDR)
+    print(msg_vgg)
+    
+    # evaluation on MUSIC
+    SISDR, SDRi = music_evaluator(pl_model)
+    msg_music = "MUSIC Avg SDRi: {:.3f}, SISDR: {:.3f}".format(SDRi, SISDR)
+    print(msg_music)
+
+    # evaluation on ESC-50
+    SISDR, SDRi = esc50_evaluator(pl_model)
+    msg_esc50 = "ESC-50 Avg SDRi: {:.3f}, SISDR: {:.3f}".format(SDRi, SISDR)
+    print(msg_esc50)
+
+    # evaluation on AudioSet
+    stats_dict = audioset_evaluator(pl_model=pl_model)
+    median_sdris = {}
+    median_sisdrs = {}
+
+    for class_id in range(527):
+        median_sdris[class_id] = np.nanmedian(stats_dict["sdris_dict"][class_id])
+        median_sisdrs[class_id] = np.nanmedian(stats_dict["sisdrs_dict"][class_id])
+
+    SDRi = get_mean_sdr_from_dict(median_sdris)
+    SISDR = get_mean_sdr_from_dict(median_sisdrs)
+    msg_audioset = "AudioSet Avg SDRi: {:.3f}, SISDR: {:.3f}".format(SDRi, SISDR)
+    print(msg_audioset)
+
+    # evaluation on AudioCaps
+    SISDR, SDRi = audiocaps_evaluator(pl_model)
+    msg_audiocaps = "AudioCaps Avg SDRi: {:.3f}, SISDR: {:.3f}".format(SDRi, SISDR)
+    print(msg_audiocaps)
+
+    # evaluation on Clotho
+    SISDR, SDRi = clotho_evaluator(pl_model)
+    msg_clotho = "Clotho Avg SDRi: {:.3f}, SISDR: {:.3f}".format(SDRi, SISDR)
+    print(msg_clotho)
+    
+    msgs = [msg_audioset, msg_vgg, msg_audiocaps, msg_clotho, msg_music, msg_esc50]
+
+    # open file in write mode
+    log_path = os.path.join(log_dir, 'eval_results.txt')
+    with open(log_path, 'w') as fp:
+        for msg in msgs:
+            fp.write(msg + '\n')
+    print(f'Eval log is written to {log_path} ...')
+    print('-------------------------  Done  ---------------------------')
+
+
+if __name__ == '__main__':
+    eval(checkpoint_path='checkpoint/audiosep_base.ckpt')
+
+   
+
+
+
+
+

callbacks/base.py

@@ -0,0 +1,35 @@
+import os
+import lightning.pytorch as pl
+from lightning.pytorch.utilities import rank_zero_only
+
+
+class CheckpointEveryNSteps(pl.Callback):
+    def __init__(
+        self,
+        checkpoints_dir,
+        save_step_frequency,
+    ) -> None:
+        r"""Save a checkpoint every N steps.
+
+        Args:
+            checkpoints_dir (str): directory to save checkpoints
+            save_step_frequency (int): save checkpoint every N step
+        """
+
+        self.checkpoints_dir = checkpoints_dir
+        self.save_step_frequency = save_step_frequency
+
+    @rank_zero_only
+    def on_train_batch_end(self, *args, **kwargs) -> None:
+        r"""Save a checkpoint every N steps."""
+
+        trainer = args[0]
+        global_step = trainer.global_step
+
+        if global_step == 1 or global_step % self.save_step_frequency == 0:
+
+            ckpt_path = os.path.join(
+                self.checkpoints_dir,
+                "step={}.ckpt".format(global_step))
+            trainer.save_checkpoint(ckpt_path)
+            print("Save checkpoint to {}".format(ckpt_path))

checkpoint/audiosep_base_4M_steps.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f8cda01bfd0ebd141eef45d41db7a3ada23a56568465840d3cff04b8010ce82c
+size 1264844076

checkpoint/music_speech_audioset_epoch_15_esc_89.98.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:51c68f12f9d7ea25fdaaccf741ec7f81e93ee594455410f3bca4f47f88d8e006
+size 2352471003

cog.yaml

@@ -0,0 +1,21 @@
+# Configuration for Cog ⚙️
+# Reference: https://github.com/replicate/cog/blob/main/docs/yaml.md
+
+build:
+  gpu: true
+  python_version: "3.11"
+  python_packages:
+    - "torchlibrosa==0.1.0"
+    - "lightning==2.1.0"
+    - "torch==2.0.1"
+    - "transformers==4.28.1"
+    - "braceexpand==0.1.7"
+    - "webdataset==0.2.60"
+    - "soundfile==0.12.1"
+    - "torchaudio==2.0.2"
+    - "torchvision==0.15.2"
+    - "h5py==3.10.0"
+    - "ftfy==6.1.1"
+    - "pandas==2.1.1"
+    - "wget==3.2"
+predict: "predict.py:Predictor"

config/audiosep_base.yaml

@@ -0,0 +1,41 @@
+---
+task_name: AudioSep
+
+data:
+    datafiles:
+        - 'datafiles/template.json'
+
+    sampling_rate: 32000
+    segment_seconds: 5
+    loudness_norm:
+        lower_db: -10
+        higher_db: 10
+    max_mix_num: 2
+
+model:
+    query_net: CLAP
+    condition_size: 512
+    model_type: ResUNet30
+    input_channels: 1
+    output_channels: 1
+    resume_checkpoint: ""
+    use_text_ratio: 1.0
+
+train:
+    optimizer:
+        optimizer_type: AdamW
+        learning_rate: 1e-3
+        warm_up_steps: 10000
+        reduce_lr_steps: 1000000
+        lr_lambda_type: constant_warm_up
+    num_nodes: 1
+    num_workers: 6
+    loss_type: l1_wav
+    sync_batchnorm: True
+    batch_size_per_device: 12
+    steps_per_epoch: 10000  # Every 10000 steps is called an `epoch`.
+    evaluate_step_frequency: 10000     # Evaluate every #evaluate_step_frequency steps.
+    save_step_frequency: 20000  # Save every #save_step_frequency steps.
+    early_stop_steps: 10000001
+    random_seed: 1234
+

data/audiotext_dataset.py

@@ -0,0 +1,91 @@
+import json
+import random
+import torch
+import torchaudio
+from torch.utils.data import Dataset
+
+
+class AudioTextDataset(Dataset):
+    """Can sample data from audio-text databases
+    Params:
+    sampling_rate: audio sampling rate
+    max_clip_len: max length (seconds) of audio clip to be sampled
+    """
+    def __init__(
+        self,
+        datafiles=[''], 
+        sampling_rate=32000, 
+        max_clip_len=5,
+    ):
+        all_data_json = []
+        for datafile in datafiles:
+            with open(datafile, 'r') as fp:
+                data_json = json.load(fp)['data']
+                all_data_json.extend(data_json)
+        self.all_data_json = all_data_json
+
+        self.sampling_rate = sampling_rate
+        self.max_length = max_clip_len * sampling_rate
+
+    def __len__(self):
+        return len(self.all_data_json)
+
+    def _cut_or_randomcrop(self, waveform):
+        # waveform: [1, samples]
+        # random crop
+        if waveform.size(1) > self.max_length:
+            random_idx = random.randint(0, waveform.size(1)-self.max_length)
+            waveform = waveform[:, random_idx:random_idx+self.max_length]
+        else:
+            temp_wav = torch.zeros(1, self.max_length)
+            temp_wav[:, 0:waveform.size(1)] = waveform
+            waveform = temp_wav
+
+        assert waveform.size(1) == self.max_length, \
+            f"number of audio samples is {waveform.size(1)}"
+
+        return waveform
+
+    def _read_audio(self, index):
+        try:
+            audio_path = self.all_data_json[index]['wav']
+            audio_data, audio_rate = torchaudio.load(audio_path, channels_first=True)
+            text = self.all_data_json[index]['caption']
+
+            # drop short utterance
+            if audio_data.size(1) < self.sampling_rate * 1:
+                raise Exception(f'{audio_path} is too short, drop it ...') 
+            
+            return text, audio_data, audio_rate
+        
+        except Exception as e:
+            print(f'error: {e} occurs, when loading {audio_path}')
+            random_index = random.randint(0, len(self.all_data_json)-1)
+            return self._read_audio(index=random_index)
+
+    def __getitem__(self, index):
+        # create a audio tensor  
+        text, audio_data, audio_rate = self._read_audio(index)
+        audio_len = audio_data.shape[1] / audio_rate
+        # convert stero to single channel
+        if audio_data.shape[0] > 1:
+            # audio_data: [samples]
+            audio_data = (audio_data[0] + audio_data[1]) / 2
+        else:
+            audio_data = audio_data.squeeze(0)
+        
+        # resample audio clip
+        if audio_rate != self.sampling_rate:
+            audio_data = torchaudio.functional.resample(audio_data, orig_freq=audio_rate, new_freq=self.sampling_rate)
+        
+        audio_data = audio_data.unsqueeze(0)
+        
+        audio_data = self._cut_or_randomcrop(audio_data)            
+
+        data_dict = {
+            'text': text, 
+            'waveform': audio_data,  
+            'modality': 'audio_text'
+        }
+
+        return data_dict

data/datamodules.py

@@ -0,0 +1,122 @@
+from typing import Dict, List, Optional, NoReturn
+import torch
+import lightning.pytorch as pl
+from torch.utils.data import DataLoader
+from data.audiotext_dataset import AudioTextDataset
+
+
+class DataModule(pl.LightningDataModule):
+    def __init__(
+        self,
+        train_dataset: object,
+        batch_size: int,
+        num_workers: int
+    ):
+        r"""Data module. To get one batch of data:
+
+        code-block:: python
+
+            data_module.setup()
+
+            for batch_data_dict in data_module.train_dataloader():
+                print(batch_data_dict.keys())
+                break
+
+        Args:
+            train_sampler: Sampler object
+            train_dataset: Dataset object
+            num_workers: int
+            distributed: bool
+        """
+        super().__init__()
+        self._train_dataset = train_dataset
+        self.num_workers = num_workers
+        self.batch_size = batch_size
+        self.collate_fn = collate_fn
+
+
+    def prepare_data(self):
+        # download, split, etc...
+        # only called on 1 GPU/TPU in distributed
+        pass
+
+    def setup(self, stage: Optional[str] = None) -> NoReturn:
+        r"""called on every device."""
+
+        # make assignments here (val/train/test split)
+        # called on every process in DDP
+
+        # SegmentSampler is used for selecting segments for training.
+        # On multiple devices, each SegmentSampler samples a part of mini-batch
+        # data.
+        self.train_dataset = self._train_dataset
+        
+        
+    def train_dataloader(self) -> torch.utils.data.DataLoader:
+        r"""Get train loader."""
+        train_loader = DataLoader(
+            dataset=self.train_dataset,
+            batch_size=self.batch_size,
+            collate_fn=self.collate_fn,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            persistent_workers=False,
+            shuffle=True
+        )
+
+        return train_loader
+
+    def val_dataloader(self):
+        # val_split = Dataset(...)
+        # return DataLoader(val_split)
+        pass
+
+    def test_dataloader(self):
+        # test_split = Dataset(...)
+        # return DataLoader(test_split)
+        pass
+
+    def teardown(self):
+        # clean up after fit or test
+        # called on every process in DDP
+        pass
+
+
+def collate_fn(list_data_dict):
+    r"""Collate mini-batch data to inputs and targets for training.
+
+    Args:
+        list_data_dict: e.g., [
+            {
+                'text': 'a sound of dog',
+                'waveform': (1, samples),
+                'modality': 'audio_text'
+            }
+            ...
+            ]
+    Returns:
+        data_dict: e.g. 
+            'audio_text': {
+                'text': ['a sound of dog', ...]
+                'waveform': (batch_size, 1, samples)
+        }
+    """
+    
+    at_list_data_dict = [data_dict for data_dict in list_data_dict if data_dict['modality']=='audio_text']
+
+    at_data_dict = {}
+    
+    if len(at_list_data_dict) > 0:
+        for key in at_list_data_dict[0].keys():
+            at_data_dict[key] = [at_data_dict[key] for at_data_dict in at_list_data_dict]
+            if key == 'waveform':
+                at_data_dict[key] = torch.stack(at_data_dict[key])
+            elif key == 'text':
+                at_data_dict[key] = [text for text in at_data_dict[key]]
+
+    
+    data_dict = {
+        'audio_text': at_data_dict
+    }
+    
+    return data_dict

data/waveform_mixers.py

@@ -0,0 +1,127 @@
+import random
+import sre_compile
+import numpy as np
+import torch
+import torch.nn as nn
+import pyloudnorm as pyln
+
+
+class SegmentMixer(nn.Module):
+    def __init__(self, max_mix_num, lower_db, higher_db):
+        super(SegmentMixer, self).__init__()
+
+        self.max_mix_num = max_mix_num
+        self.loudness_param = {
+            'lower_db': lower_db,
+            'higher_db': higher_db,
+        }
+
+    def __call__(self, waveforms):
+        
+        batch_size = waveforms.shape[0]
+
+        data_dict = {
+            'segment': [],
+            'mixture': [],
+        }
+
+        for n in range(0, batch_size):
+
+            segment = waveforms[n].clone()
+
+            # create zero tensors as the background template
+            noise = torch.zeros_like(segment)
+
+            mix_num = random.randint(2, self.max_mix_num)
+            assert mix_num >= 2
+
+            for i in range(1, mix_num):
+                next_segment = waveforms[(n + i) % batch_size]
+                rescaled_next_segment = dynamic_loudnorm(audio=next_segment, reference=segment, **self.loudness_param)
+                noise += rescaled_next_segment
+
+            # randomly normalize background noise
+            noise = dynamic_loudnorm(audio=noise, reference=segment, **self.loudness_param)
+
+            # create audio mixyure
+            mixture = segment + noise
+
+            # declipping if need be
+            max_value = torch.max(torch.abs(mixture))
+            if max_value > 1:
+                segment *= 0.9 / max_value
+                mixture *= 0.9 / max_value
+
+            data_dict['segment'].append(segment)
+            data_dict['mixture'].append(mixture)
+
+        for key in data_dict.keys():
+            data_dict[key] = torch.stack(data_dict[key], dim=0)
+
+        # return data_dict
+        return data_dict['mixture'], data_dict['segment']
+
+
+def rescale_to_match_energy(segment1, segment2):
+
+    ratio = get_energy_ratio(segment1, segment2)
+    rescaled_segment1 = segment1 / ratio
+    return rescaled_segment1 
+
+
+def get_energy(x):
+    return torch.mean(x ** 2)
+
+
+def get_energy_ratio(segment1, segment2):
+
+    energy1 = get_energy(segment1)
+    energy2 = max(get_energy(segment2), 1e-10)
+    ratio = (energy1 / energy2) ** 0.5
+    ratio = torch.clamp(ratio, 0.02, 50)
+    return ratio
+
+
+def dynamic_loudnorm(audio, reference, lower_db=-10, higher_db=10): 
+    rescaled_audio = rescale_to_match_energy(audio, reference)
+    
+    delta_loudness = random.randint(lower_db, higher_db)
+
+    gain = np.power(10.0, delta_loudness / 20.0)
+
+    return gain * rescaled_audio
+
+
+def torch_to_numpy(tensor):
+    """Convert a PyTorch tensor to a NumPy array."""
+    if isinstance(tensor, torch.Tensor):
+        return tensor.detach().cpu().numpy()
+    else:
+        raise ValueError("Input must be a PyTorch tensor.")
+
+
+def numpy_to_torch(array):
+    """Convert a NumPy array to a PyTorch tensor."""
+    if isinstance(array, np.ndarray):
+        return torch.from_numpy(array)
+    else:
+        raise ValueError("Input must be a NumPy array.")
+
+
+# decayed
+def random_loudness_norm(audio, lower_db=-35, higher_db=-15, sr=32000):
+    device = audio.device
+    audio = torch_to_numpy(audio.squeeze(0))
+    # randomly select a norm volume
+    norm_vol = random.randint(lower_db, higher_db)
+
+    # measure the loudness first 
+    meter = pyln.Meter(sr) # create BS.1770 meter
+    loudness = meter.integrated_loudness(audio)
+    # loudness normalize audio
+    normalized_audio = pyln.normalize.loudness(audio, loudness, norm_vol)
+
+    normalized_audio = numpy_to_torch(normalized_audio).unsqueeze(0)
+    
+    return normalized_audio.to(device)
+    

datafiles/template.json

@@ -0,0 +1,8 @@
+{
+    "data": [
+     {
+      "wav": "path_to_audio_file",
+      "caption": "textual_desciptions"
+     }
+    ]
+}

environment.yml

@@ -0,0 +1,326 @@
+name: AudioSep
+channels:
+  - pytorch
+  - nvidia
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - backcall=0.2.0=pyhd3eb1b0_0
+  - blas=1.0=mkl
+  - boltons=23.0.0=py310h06a4308_0
+  - brotlipy=0.7.0=py310h7f8727e_1002
+  - bzip2=1.0.8=h7b6447c_0
+  - ca-certificates=2023.01.10=h06a4308_0
+  - certifi=2022.12.7=py310h06a4308_0
+  - cffi=1.15.1=py310h5eee18b_3
+  - charset-normalizer=2.0.4=pyhd3eb1b0_0
+  - comm=0.1.2=py310h06a4308_0
+  - conda=23.3.1=py310h06a4308_0
+  - conda-content-trust=0.1.3=py310h06a4308_0
+  - conda-package-handling=2.0.2=py310h06a4308_0
+  - conda-package-streaming=0.7.0=py310h06a4308_0
+  - cryptography=38.0.4=py310h9ce1e76_0
+  - cuda=11.6.1=0
+  - cuda-cccl=11.6.55=hf6102b2_0
+  - cuda-command-line-tools=11.6.2=0
+  - cuda-compiler=11.6.2=0
+  - cuda-cudart=11.6.55=he381448_0
+  - cuda-cudart-dev=11.6.55=h42ad0f4_0
+  - cuda-cuobjdump=11.6.124=h2eeebcb_0
+  - cuda-cupti=11.6.124=h86345e5_0
+  - cuda-cuxxfilt=11.6.124=hecbf4f6_0
+  - cuda-driver-dev=11.6.55=0
+  - cuda-gdb=12.1.55=0
+  - cuda-libraries=11.6.1=0
+  - cuda-libraries-dev=11.6.1=0
+  - cuda-memcheck=11.8.86=0
+  - cuda-nsight=12.1.55=0
+  - cuda-nsight-compute=12.1.0=0
+  - cuda-nvcc=11.6.124=hbba6d2d_0
+  - cuda-nvdisasm=12.1.55=0
+  - cuda-nvml-dev=11.6.55=haa9ef22_0
+  - cuda-nvprof=12.1.55=0
+  - cuda-nvprune=11.6.124=he22ec0a_0
+  - cuda-nvrtc=11.6.124=h020bade_0
+  - cuda-nvrtc-dev=11.6.124=h249d397_0
+  - cuda-nvtx=11.6.124=h0630a44_0
+  - cuda-nvvp=12.1.55=0
+  - cuda-runtime=11.6.1=0
+  - cuda-samples=11.6.101=h8efea70_0
+  - cuda-sanitizer-api=12.1.55=0
+  - cuda-toolkit=11.6.1=0
+  - cuda-tools=11.6.1=0
+  - cuda-visual-tools=11.6.1=0
+  - debugpy=1.5.1=py310h295c915_0
+  - decorator=5.1.1=pyhd3eb1b0_0
+  - flit-core=3.8.0=py310h06a4308_0
+  - freetype=2.12.1=h4a9f257_0
+  - gds-tools=1.6.0.25=0
+  - giflib=5.2.1=h5eee18b_3
+  - gmp=6.2.1=h295c915_3
+  - gnutls=3.6.15=he1e5248_0
+  - idna=3.4=py310h06a4308_0
+  - intel-openmp=2021.4.0=h06a4308_3561
+  - ipykernel=6.19.2=py310h2f386ee_0
+  - ipython=8.12.0=py310h06a4308_0
+  - jpeg=9e=h5eee18b_1
+  - jsonpatch=1.32=pyhd3eb1b0_0
+  - jsonpointer=2.1=pyhd3eb1b0_0
+  - jupyter_client=8.1.0=py310h06a4308_0
+  - jupyter_core=5.3.0=py310h06a4308_0
+  - lame=3.100=h7b6447c_0
+  - lcms2=2.12=h3be6417_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - lerc=3.0=h295c915_0
+  - libcublas=11.9.2.110=h5e84587_0
+  - libcublas-dev=11.9.2.110=h5c901ab_0
+  - libcufft=10.7.1.112=hf425ae0_0
+  - libcufft-dev=10.7.1.112=ha5ce4c0_0
+  - libcufile=1.6.0.25=0
+  - libcufile-dev=1.6.0.25=0
+  - libcurand=10.3.2.56=0
+  - libcurand-dev=10.3.2.56=0
+  - libcusolver=11.3.4.124=h33c3c4e_0
+  - libcusparse=11.7.2.124=h7538f96_0
+  - libcusparse-dev=11.7.2.124=hbbe9722_0
+  - libdeflate=1.17=h5eee18b_0
+  - libffi=3.4.2=h6a678d5_6
+  - libgcc-ng=11.2.0=h1234567_1
+  - libgomp=11.2.0=h1234567_1
+  - libiconv=1.16=h7f8727e_2
+  - libidn2=2.3.2=h7f8727e_0
+  - libnpp=11.6.3.124=hd2722f0_0
+  - libnpp-dev=11.6.3.124=h3c42840_0
+  - libnvjpeg=11.6.2.124=hd473ad6_0
+  - libnvjpeg-dev=11.6.2.124=hb5906b9_0
+  - libpng=1.6.39=h5eee18b_0
+  - libsodium=1.0.18=h7b6447c_0
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - libtasn1=4.19.0=h5eee18b_0
+  - libtiff=4.5.0=h6a678d5_2
+  - libunistring=0.9.10=h27cfd23_0
+  - libuuid=1.41.5=h5eee18b_0
+  - libwebp=1.2.4=h11a3e52_1
+  - libwebp-base=1.2.4=h5eee18b_1
+  - lz4-c=1.9.4=h6a678d5_0
+  - matplotlib-inline=0.1.6=py310h06a4308_0
+  - mkl=2021.4.0=h06a4308_640
+  - mkl-service=2.4.0=py310h7f8727e_0
+  - mkl_fft=1.3.1=py310hd6ae3a3_0
+  - mkl_random=1.2.2=py310h00e6091_0
+  - ncurses=6.4=h6a678d5_0
+  - nest-asyncio=1.5.6=py310h06a4308_0
+  - nettle=3.7.3=hbbd107a_1
+  - nsight-compute=2023.1.0.15=0
+  - numpy=1.23.5=py310hd5efca6_0
+  - numpy-base=1.23.5=py310h8e6c178_0
+  - openh264=2.1.1=h4ff587b_0
+  - openssl=1.1.1t=h7f8727e_0
+  - packaging=23.0=py310h06a4308_0
+  - parso=0.8.3=pyhd3eb1b0_0
+  - pexpect=4.8.0=pyhd3eb1b0_3
+  - pickleshare=0.7.5=pyhd3eb1b0_1003
+  - pip=22.3.1=py310h06a4308_0
+  - platformdirs=2.5.2=py310h06a4308_0
+  - pluggy=1.0.0=py310h06a4308_1
+  - psutil=5.9.0=py310h5eee18b_0
+  - ptyprocess=0.7.0=pyhd3eb1b0_2
+  - pure_eval=0.2.2=pyhd3eb1b0_0
+  - pycosat=0.6.4=py310h5eee18b_0
+  - pycparser=2.21=pyhd3eb1b0_0
+  - pyopenssl=22.0.0=pyhd3eb1b0_0
+  - pysocks=1.7.1=py310h06a4308_0
+  - python=3.10.9=h7a1cb2a_0
+  - python-dateutil=2.8.2=pyhd3eb1b0_0
+  - pytorch=1.13.1=py3.10_cuda11.6_cudnn8.3.2_0
+  - pytorch-cuda=11.6=h867d48c_1
+  - pytorch-mutex=1.0=cuda
+  - pyzmq=23.2.0=py310h6a678d5_0
+  - readline=8.2=h5eee18b_0
+  - requests=2.28.1=py310h06a4308_0
+  - ruamel.yaml=0.17.21=py310h5eee18b_0
+  - ruamel.yaml.clib=0.2.6=py310h5eee18b_1
+  - setuptools=65.6.3=py310h06a4308_0
+  - six=1.16.0=pyhd3eb1b0_1
+  - sqlite=3.40.1=h5082296_0
+  - stack_data=0.2.0=pyhd3eb1b0_0
+  - tk=8.6.12=h1ccaba5_0
+  - toolz=0.12.0=py310h06a4308_0
+  - torchaudio=0.13.1=py310_cu116
+  - torchvision=0.14.1=py310_cu116
+  - tornado=6.2=py310h5eee18b_0
+  - tqdm=4.64.1=py310h06a4308_0
+  - typing_extensions=4.4.0=py310h06a4308_0
+  - tzdata=2022g=h04d1e81_0
+  - urllib3=1.26.14=py310h06a4308_0
+  - wheel=0.37.1=pyhd3eb1b0_0
+  - xz=5.2.10=h5eee18b_1
+  - zeromq=4.3.4=h2531618_0
+  - zlib=1.2.13=h5eee18b_0
+  - zstandard=0.18.0=py310h5eee18b_0
+  - zstd=1.5.4=hc292b87_0
+  - pip:
+      - absl-py==1.4.0
+      - aiohttp==3.8.4
+      - aiosignal==1.3.1
+      - anyio==3.6.2
+      - appdirs==1.4.4
+      - arrow==1.2.3
+      - asttokens==2.2.1
+      - async-generator==1.10
+      - async-timeout==4.0.2
+      - attrs==22.2.0
+      - audioread==3.0.0
+      - av==10.0.0
+      - beartype==0.12.0
+      - beautifulsoup4==4.12.2
+      - blessed==1.20.0
+      - braceexpand==0.1.7
+      - cachetools==5.3.0
+      - click==8.1.3
+      - contourpy==1.0.7
+      - croniter==1.3.10
+      - cycler==0.11.0
+      - dataclasses-json==0.5.8
+      - dateutils==0.6.12
+      - decord==0.6.0
+      - deepdiff==6.3.0
+      - dtk==0.2
+      - exceptiongroup==1.1.1
+      - executing==1.2.0
+      - fastapi==0.88.0
+      - ffmpeg==1.4
+      - ffmpeg-python==0.2.0
+      - filelock==3.12.0
+      - fonttools==4.39.3
+      - frozenlist==1.3.3
+      - fsspec==2023.4.0
+      - ftfy==6.1.1
+      - future==0.18.3
+      - gammatone==1.0
+      - google-auth==2.17.3
+      - google-auth-oauthlib==1.0.0
+      - greenlet==2.0.2
+      - grpcio==1.54.0
+      - h11==0.14.0
+      - h5py==3.8.0
+      - hickle==5.0.2
+      - huggingface-hub==0.14.1
+      - humanize==4.6.0
+      - imageio==2.27.0
+      - inquirer==3.1.3
+      - ipdb==0.13.13
+      - itsdangerous==2.1.2
+      - jedi==0.18.2
+      - jinja2==3.1.2
+      - joblib==1.2.0
+      - kiwisolver==1.4.4
+      - langchain==0.0.216
+      - langchainplus-sdk==0.0.17
+      - lazy-loader==0.2
+      - librosa==0.10.0.post2
+      - lightning==2.0.0
+      - lightning-cloud==0.5.33
+      - lightning-utilities==0.8.0
+      - llvmlite==0.39.1
+      - markdown==3.4.3
+      - markdown-it-py==2.2.0
+      - markupsafe==2.1.2
+      - marshmallow==3.19.0
+      - marshmallow-enum==1.5.1
+      - matplotlib==3.7.1
+      - mdurl==0.1.2
+      - mergedeep==1.3.4
+      - mock==5.0.2
+      - msgpack==1.0.5
+      - msgpack-numpy==0.4.8
+      - multidict==6.0.4
+      - musdb==0.4.0
+      - mypy-extensions==1.0.0
+      - networkx==3.1
+      - nose==1.3.7
+      - numba==0.56.4
+      - numexpr==2.8.4
+      - oauthlib==3.2.2
+      - openai==0.27.8
+      - openapi-schema-pydantic==1.2.4
+      - opencv-python==4.7.0.72
+      - ordered-set==4.1.0
+      - outcome==1.2.0
+      - pandas==1.5.3
+      - panns-inference==0.1.0
+      - pesq==0.0.4
+      - pillow==9.5.0
+      - pooch==1.6.0
+      - prompt-toolkit==3.0.38
+      - protobuf==4.22.3
+      - pyaml==23.5.9
+      - pyasn1==0.5.0
+      - pyasn1-modules==0.3.0
+      - pydantic==1.10.7
+      - pygments==2.14.0
+      - pyjwt==2.6.0
+      - pyloudnorm==0.1.1
+      - pyparsing==3.0.9
+      - pystoi==0.3.3
+      - python-editor==1.0.4
+      - python-multipart==0.0.6
+      - pytorch-ignite==0.3.0
+      - pytorch-lightning==2.0.1.post0
+      - pytz==2023.3
+      - pywavelets==1.4.1
+      - pyyaml==6.0
+      - readchar==4.0.5
+      - regex==2023.3.23
+      - requests-oauthlib==1.3.1
+      - resampy==0.4.2
+      - rich==13.3.3
+      - rsa==4.9
+      - scikit-image==0.20.0
+      - scikit-learn==1.2.2
+      - scipy==1.10.1
+      - selenium==4.8.3
+      - simplejpeg==1.6.6
+      - sniffio==1.3.0
+      - sortedcontainers==2.4.0
+      - soundfile==0.12.1
+      - soupsieve==2.4
+      - soxr==0.3.5
+      - sqlalchemy==2.0.17
+      - stack-data==0.6.2
+      - starlette==0.22.0
+      - starsessions==1.3.0
+      - stempeg==0.2.3
+      - tenacity==8.2.2
+      - tensorboard==2.12.2
+      - tensorboard-data-server==0.7.0
+      - tensorboard-plugin-wit==1.8.1
+      - termcolor==1.1.0
+      - threadpoolctl==3.1.0
+      - tifffile==2023.3.21
+      - timm==0.3.2
+      - tokenizers==0.13.3
+      - tomli==2.0.1
+      - torchfile==0.1.0
+      - torchlibrosa==0.1.0
+      - torchmetrics==0.11.4
+      - traitlets==5.9.0
+      - transformers==4.28.1
+      - trio==0.22.0
+      - trio-websocket==0.10.2
+      - typeguard==3.0.2
+      - typing-extensions==4.5.0
+      - typing-inspect==0.9.0
+      - uvicorn==0.21.1
+      - visdom==0.1.8.9
+      - wcwidth==0.2.6
+      - webdataset==0.2.48
+      - websocket-client==1.5.1
+      - websockets==11.0.1
+      - werkzeug==2.2.3
+      - wget==3.2
+      - wsproto==1.2.0
+      - yarl==1.8.2
+      - zenodo-get==1.3.4
+      - zsvision==0.7.8

evaluation/evaluate_audiocaps.py

@@ -0,0 +1,110 @@
+import os
+import sys
+import re
+from typing import Dict, List
+
+import csv
+import pandas as pd
+import numpy as np
+import torch
+from tqdm import tqdm
+import pathlib
+import librosa
+import lightning.pytorch as pl
+from models.clap_encoder import CLAP_Encoder
+
+sys.path.append('../AudioSep/')
+from utils import (
+    load_ss_model,
+    calculate_sdr,
+    calculate_sisdr,
+    parse_yaml,
+    get_mean_sdr_from_dict,
+)
+
+
+class AudioCapsEvaluator:
+    def __init__(
+        self,
+        query='caption',
+        sampling_rate=32000,
+    ) -> None:
+        r"""AudioCaps evaluator.
+
+        Args:
+            query (str): type of query, 'caption' or 'labels'
+        Returns:
+            None
+        """
+
+        self.query = query
+        self.sampling_rate = sampling_rate
+
+        with open(f'evaluation/metadata/audiocaps_eval.csv') as csv_file:
+            csv_reader = csv.reader(csv_file, delimiter=',')
+            eval_list = [row for row in csv_reader][1:]
+        
+        self.eval_list = eval_list
+        self.audio_dir = f'evaluation/data/audiocaps'
+
+    def __call__(
+        self,
+        pl_model: pl.LightningModule
+    ) -> Dict:
+        r"""Evalute."""
+
+        print(f'Evaluation on AudioCaps with [{self.query}] queries.')
+        
+        pl_model.eval()
+        device = pl_model.device
+
+        sisdrs_list = []
+        sdris_list = []
+
+        with torch.no_grad():
+            for eval_data in tqdm(self.eval_list):
+
+                idx, caption, labels, _, _ = eval_data
+
+                source_path = os.path.join(self.audio_dir, f'segment-{idx}.wav')
+                mixture_path = os.path.join(self.audio_dir, f'mixture-{idx}.wav')
+
+                source, fs = librosa.load(source_path, sr=self.sampling_rate, mono=True)
+                mixture, fs = librosa.load(mixture_path, sr=self.sampling_rate, mono=True)
+
+                sdr_no_sep = calculate_sdr(ref=source, est=mixture)
+                                
+                if self.query == 'caption':
+                    text = [caption]
+                elif self.query == 'labels':
+                    text = [labels]
+
+                conditions = pl_model.query_encoder.get_query_embed(
+                    modality='text',
+                    text=text,
+                    device=device 
+                )
+                    
+                input_dict = {
+                    "mixture": torch.Tensor(mixture)[None, None, :].to(device),
+                    "condition": conditions,
+                }
+                
+                
+                sep_segment = pl_model.ss_model(input_dict)["waveform"]
+                    # sep_segment: (batch_size=1, channels_num=1, segment_samples)
+
+                sep_segment = sep_segment.squeeze(0).squeeze(0).data.cpu().numpy()
+                    # sep_segment: (segment_samples,)
+
+                sdr = calculate_sdr(ref=source, est=sep_segment)
+                sdri = sdr - sdr_no_sep
+                sisdr = calculate_sisdr(ref=source, est=sep_segment)
+
+                sisdrs_list.append(sisdr)
+                sdris_list.append(sdri)
+
+        mean_sisdr = np.mean(sisdrs_list)
+        mean_sdri = np.mean(sdris_list)
+        
+        return mean_sisdr, mean_sdri

evaluation/evaluate_audioset.py

@@ -0,0 +1,155 @@
+import os
+import sys
+import re
+from typing import Dict, List
+
+import pandas as pd
+import numpy as np
+import torch
+from tqdm import tqdm
+import pathlib
+import librosa
+import lightning.pytorch as pl
+from models.clap_encoder import CLAP_Encoder
+
+sys.path.append('../AudioSep/')
+from utils import (
+    load_ss_model,
+    calculate_sdr,
+    calculate_sisdr,
+    parse_yaml,
+    get_mean_sdr_from_dict,
+)
+
+
+meta_csv_file = "evaluation/metadata/class_labels_indices.csv"
+df = pd.read_csv(meta_csv_file, sep=',')
+
+IDS = df['mid'].tolist()
+LABELS = df['display_name'].tolist()
+
+CLASSES_NUM = len(LABELS)
+
+IX_TO_LB = {i : label for i, label in enumerate(LABELS)}
+
+
+class AudioSetEvaluator:
+    def __init__(
+        self,
+        audios_dir='evaluation/data/audioset',
+        classes_num=527,
+        sampling_rate=32000,
+        number_per_class=10,
+    ) -> None:
+        r"""AudioSet evaluator.
+
+        Args:
+            audios_dir (str): directory of evaluation segments
+            classes_num (int): the number of sound classes
+            number_per_class (int), the number of samples to evaluate for each sound class
+
+        Returns:
+            None
+        """
+
+        self.audios_dir = audios_dir
+        self.classes_num = classes_num
+        self.number_per_class = number_per_class
+        self.sampling_rate = sampling_rate
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        pl_model: pl.LightningModule
+    ) -> Dict:
+        r"""Evalute."""
+
+        pl_model.eval()
+
+        sisdrs_dict = {class_id: [] for class_id in range(self.classes_num)}
+        sdris_dict = {class_id: [] for class_id in range(self.classes_num)}
+        
+        print('Evaluation on AudioSet with [text label] queries.')
+        
+        for class_id in tqdm(range(self.classes_num)):
+
+            sub_dir = os.path.join(
+                self.audios_dir,
+                "class_id={}".format(class_id))
+
+            audio_names = self._get_audio_names(audios_dir=sub_dir)
+
+            for audio_index, audio_name in enumerate(audio_names):
+
+                if audio_index == self.number_per_class:
+                    break
+
+                source_path = os.path.join(
+                    sub_dir, "{},source.wav".format(audio_name))
+                mixture_path = os.path.join(
+                    sub_dir, "{},mixture.wav".format(audio_name))
+
+                source, fs = librosa.load(source_path, sr=self.sampling_rate, mono=True)
+                mixture, fs = librosa.load(mixture_path, sr=self.sampling_rate, mono=True)
+
+                sdr_no_sep = calculate_sdr(ref=source, est=mixture)
+
+                device = pl_model.device
+
+                text = [IX_TO_LB[class_id]]
+
+                conditions = pl_model.query_encoder.get_query_embed(
+                    modality='text',
+                    text=text,
+                    device=device 
+                )
+
+                input_dict = {
+                    "mixture": torch.Tensor(mixture)[None, None, :].to(device),
+                    "condition": conditions,
+                }
+
+                sep_segment = pl_model.ss_model(input_dict)["waveform"]
+                # sep_segment: (batch_size=1, channels_num=1, segment_samples)
+
+                sep_segment = sep_segment.squeeze(0).squeeze(0).data.cpu().numpy()
+                # sep_segment: (segment_samples,)
+
+                sdr = calculate_sdr(ref=source, est=sep_segment)
+                sdri = sdr - sdr_no_sep
+                sisdr = calculate_sisdr(ref=source, est=sep_segment)
+
+
+                sisdrs_dict[class_id].append(sisdr)
+                sdris_dict[class_id].append(sdri)
+
+
+        stats_dict = {
+            "sisdrs_dict": sisdrs_dict,
+            "sdris_dict": sdris_dict,
+        }
+
+        return stats_dict
+
+    def _get_audio_names(self, audios_dir: str) -> List[str]:
+        r"""Get evaluation audio names."""
+        audio_names = sorted(os.listdir(audios_dir))
+
+        audio_names = [audio_name for audio_name in audio_names if '.wav' in audio_name]
+        
+        audio_names = [
+            re.search(
+                "(.*),(mixture|source).wav",
+                audio_name).group(1) for audio_name in audio_names]
+
+        audio_names = sorted(list(set(audio_names)))
+
+        return audio_names
+
+    @staticmethod
+    def get_median_metrics(stats_dict, metric_type):
+        class_ids = stats_dict[metric_type].keys()
+        median_stats_dict = {
+            class_id: np.nanmedian(
+                stats_dict[metric_type][class_id]) for class_id in class_ids}
+        return median_stats_dict

evaluation/evaluate_clotho.py

@@ -0,0 +1,102 @@
+import os
+import sys
+import re
+from typing import Dict, List
+
+import csv
+import pandas as pd
+import numpy as np
+import torch
+from tqdm import tqdm
+import pathlib
+import librosa
+import lightning.pytorch as pl
+from models.clap_encoder import CLAP_Encoder
+
+sys.path.append('../AudioSep/')
+from utils import (
+    load_ss_model,
+    calculate_sdr,
+    calculate_sisdr,
+    parse_yaml,
+    get_mean_sdr_from_dict,
+)
+
+
+class ClothoEvaluator:
+    def __init__(
+        self,
+        sampling_rate=32000,
+    ) -> None:
+        r"""Clotho evaluator.
+        Returns:
+            None
+        """
+
+        self.sampling_rate = sampling_rate
+
+        with open('evaluation/metadata/clotho_eval.csv') as csv_file:
+            csv_reader = csv.reader(csv_file, delimiter=',')
+            eval_list = [row for row in csv_reader][1:]
+        
+        self.eval_list = eval_list
+        self.audio_dir = 'evaluation/data/clotho'
+
+    def __call__(
+        self,
+        pl_model: pl.LightningModule
+    ) -> Dict:
+        r"""Evalute."""
+
+        print(f'Evaluation on Clotho Evaluation with [caption] queries.')
+        
+        pl_model.eval()
+        device = pl_model.device
+
+        sisdrs_list = []
+        sdris_list = []
+
+        with torch.no_grad():
+            for eval_data in tqdm(self.eval_list):
+
+                idx, caption, _, _, _ = eval_data
+
+                source_path = os.path.join(self.audio_dir, f'segment-{idx}.wav')
+                mixture_path = os.path.join(self.audio_dir, f'mixture-{idx}.wav')
+
+                source, fs = librosa.load(source_path, sr=self.sampling_rate, mono=True)
+                mixture, fs = librosa.load(mixture_path, sr=self.sampling_rate, mono=True)
+
+                sdr_no_sep = calculate_sdr(ref=source, est=mixture)
+                                
+                text = [caption]
+
+                conditions = pl_model.query_encoder.get_query_embed(
+                    modality='text',
+                    text=text,
+                    device=device 
+                )
+
+                input_dict = {
+                    "mixture": torch.Tensor(mixture)[None, None, :].to(device),
+                    "condition": conditions,
+                }
+
+                sep_segment = pl_model.ss_model(input_dict)["waveform"]
+                # sep_segment: (batch_size=1, channels_num=1, segment_samples)
+
+                sep_segment = sep_segment.squeeze(0).squeeze(0).data.cpu().numpy()
+                # sep_segment: (segment_samples,)
+
+                sdr = calculate_sdr(ref=source, est=sep_segment)
+                sdri = sdr - sdr_no_sep
+                sisdr = calculate_sisdr(ref=source, est=sep_segment)
+
+
+                sisdrs_list.append(sisdr)
+                sdris_list.append(sdri)
+
+        mean_sisdr = np.mean(sisdrs_list)
+        mean_sdri = np.mean(sdris_list)
+        
+        return mean_sisdr, mean_sdri

evaluation/evaluate_esc50.py

@@ -0,0 +1,102 @@
+import os
+import sys
+import re
+from typing import Dict, List
+
+import csv
+import pandas as pd
+import numpy as np
+import torch
+from tqdm import tqdm
+import pathlib
+import librosa
+import lightning.pytorch as pl
+from models.clap_encoder import CLAP_Encoder
+
+sys.path.append('../AudioSep/')
+from utils import (
+    load_ss_model,
+    calculate_sdr,
+    calculate_sisdr,
+    parse_yaml,
+    get_mean_sdr_from_dict,
+)
+
+
+class ESC50Evaluator:
+    def __init__(
+        self,
+        sampling_rate=32000
+    ) -> None:
+        r"""ESC-50 evaluator.
+
+        Returns:
+            None
+        """
+
+        self.sampling_rate = sampling_rate
+
+        with open('evaluation/metadata/esc50_eval.csv') as csv_file:
+            csv_reader = csv.reader(csv_file, delimiter=',')
+            eval_list = [row for row in csv_reader][1:]
+        
+        self.eval_list = eval_list
+        self.audio_dir = 'evaluation/data/esc50'
+
+    def __call__(
+        self,
+        pl_model: pl.LightningModule
+    ) -> Dict:
+        r"""Evalute."""
+
+        print(f'Evaluation on ESC-50 with [text label] queries.')
+        
+        pl_model.eval()
+        device = pl_model.device
+
+        sisdrs_list = []
+        sdris_list = []
+
+        with torch.no_grad():
+            for eval_data in tqdm(self.eval_list):
+
+                idx, caption, _, _, = eval_data
+
+                source_path = os.path.join(self.audio_dir, f'segment-{idx}.wav')
+                mixture_path = os.path.join(self.audio_dir, f'mixture-{idx}.wav')
+
+                source, fs = librosa.load(source_path, sr=self.sampling_rate, mono=True)
+                mixture, fs = librosa.load(mixture_path, sr=self.sampling_rate, mono=True)
+
+                sdr_no_sep = calculate_sdr(ref=source, est=mixture)
+                                
+                text = [caption]
+
+                conditions = pl_model.query_encoder.get_query_embed(
+                    modality='text',
+                    text=text,
+                    device=device 
+                )
+                    
+                input_dict = {
+                    "mixture": torch.Tensor(mixture)[None, None, :].to(device),
+                    "condition": conditions,
+                }
+                
+                sep_segment = pl_model.ss_model(input_dict)["waveform"]
+                    # sep_segment: (batch_size=1, channels_num=1, segment_samples)
+
+                sep_segment = sep_segment.squeeze(0).squeeze(0).data.cpu().numpy()
+                    # sep_segment: (segment_samples,)
+                
+                sdr = calculate_sdr(ref=source, est=sep_segment)
+                sdri = sdr - sdr_no_sep
+                sisdr = calculate_sisdr(ref=source, est=sep_segment)
+
+                sisdrs_list.append(sisdr)
+                sdris_list.append(sdri)
+        
+        mean_sdri = np.mean(sdris_list)
+        mean_sisdr = np.mean(sisdrs_list)
+        
+        return mean_sisdr, mean_sdri

evaluation/evaluate_music.py

@@ -0,0 +1,118 @@
+import os
+import sys
+import re
+from typing import Dict, List
+
+import csv
+import pandas as pd
+import numpy as np
+import torch
+from tqdm import tqdm
+import pathlib
+import librosa
+import lightning.pytorch as pl
+from models.clap_encoder import CLAP_Encoder
+
+sys.path.append('../AudioSep/')
+from utils import (
+    load_ss_model,
+    calculate_sdr,
+    calculate_sisdr,
+    parse_yaml,
+    get_mean_sdr_from_dict,
+)
+
+
+class MUSICEvaluator:
+    def __init__(
+        self,
+        sampling_rate=32000
+    ) -> None:
+
+        self.sampling_rate = sampling_rate
+
+        with open('evaluation/metadata/music_eval.csv') as csv_file:
+            csv_reader = csv.reader(csv_file, delimiter=',')
+            eval_list = [row for row in csv_reader][1:]
+        
+        self.eval_list = eval_list
+        self.audio_dir = 'evaluation/data/music'
+
+        self.source_types = [
+        "acoustic guitar", 
+        "violin", 
+        "accordion", 
+        "xylophone", 
+        "erhu", 
+        "trumpet", 
+        "tuba", 
+        "cello", 
+        "flute", 
+        "saxophone"]
+
+    def __call__(
+        self,
+        pl_model: pl.LightningModule
+    ) -> Dict:
+        r"""Evalute."""
+
+        print(f'Evaluation on MUSIC Test with [text label] queries.')
+        
+        pl_model.eval()
+        device = pl_model.device
+
+        sisdrs_list = {source_type: [] for source_type in self.source_types}
+        sdris_list = {source_type: [] for source_type in self.source_types}
+
+        with torch.no_grad():
+            for eval_data in tqdm(self.eval_list):
+
+                idx, caption, _, _, = eval_data
+
+                source_path = os.path.join(self.audio_dir, f'segment-{idx}.wav')
+                mixture_path = os.path.join(self.audio_dir, f'mixture-{idx}.wav')
+
+                source, fs = librosa.load(source_path, sr=self.sampling_rate, mono=True)
+                mixture, fs = librosa.load(mixture_path, sr=self.sampling_rate, mono=True)
+
+                sdr_no_sep = calculate_sdr(ref=source, est=mixture)
+                                
+                text = [caption]
+
+                conditions = pl_model.query_encoder.get_query_embed(
+                    modality='text',
+                    text=text,
+                    device=device 
+                )
+                    
+                input_dict = {
+                    "mixture": torch.Tensor(mixture)[None, None, :].to(device),
+                    "condition": conditions,
+                }
+                
+                sep_segment = pl_model.ss_model(input_dict)["waveform"]
+                    # sep_segment: (batch_size=1, channels_num=1, segment_samples)
+
+                sep_segment = sep_segment.squeeze(0).squeeze(0).data.cpu().numpy()
+                    # sep_segment: (segment_samples,)
+
+                sdr = calculate_sdr(ref=source, est=sep_segment)
+                sdri = sdr - sdr_no_sep
+                sisdr = calculate_sisdr(ref=source, est=sep_segment)
+
+                sisdrs_list[caption].append(sisdr)
+                sdris_list[caption].append(sdri)
+
+        mean_sisdr_list = []
+        mean_sdri_list = []
+        
+        for source_class in self.source_types:
+            sisdr = np.mean(sisdrs_list[source_class])
+            sdri = np.mean(sdris_list[source_class])
+            mean_sisdr_list.append(sisdr)
+            mean_sdri_list.append(sdri)
+        
+        mean_sdri = np.mean(mean_sdri_list)
+        mean_sisdr = np.mean(mean_sisdr_list)
+        
+        return mean_sisdr, mean_sdri

evaluation/evaluate_vggsound.py

@@ -0,0 +1,114 @@
+import os
+import sys
+import re
+from typing import Dict, List
+
+import csv
+import pandas as pd
+import numpy as np
+import torch
+from tqdm import tqdm
+import pathlib
+import librosa
+import lightning.pytorch as pl
+from models.clap_encoder import CLAP_Encoder
+
+sys.path.append('../AudioSep/')
+from utils import (
+    load_ss_model,
+    calculate_sdr,
+    calculate_sisdr,
+    parse_yaml,
+    get_mean_sdr_from_dict,
+)
+
+
+class VGGSoundEvaluator:
+    def __init__(
+        self,
+        sampling_rate=32000
+    ) -> None:
+        r"""VGGSound evaluator.
+
+        Args:
+            data_recipe (str): dataset split, 'yan' 
+        Returns:
+            None
+        """
+
+        self.sampling_rate = sampling_rate
+
+        with open('evaluation/metadata/vggsound_eval.csv') as csv_file:
+            csv_reader = csv.reader(csv_file, delimiter=',')
+            eval_list = [row for row in csv_reader][1:]
+        
+        self.eval_list = eval_list
+        self.audio_dir = 'evaluation/data/vggsound'
+
+    def __call__(
+        self,
+        pl_model: pl.LightningModule
+    ) -> Dict:
+        r"""Evalute."""
+
+        print(f'Evaluation on VGGSound+ with [text label] queries.')
+        
+        pl_model.eval()
+        device = pl_model.device
+
+        sisdrs_list = []
+        sdris_list = []
+        sisdris_list = []
+        
+
+        with torch.no_grad():
+            for eval_data in tqdm(self.eval_list):
+
+                # labels, source_path, mixture_path = eval_data
+                file_id, mix_wav, s0_wav, s0_text, s1_wav, s1_text = eval_data
+
+                labels = s0_text
+
+                mixture_path = os.path.join(self.audio_dir, mix_wav)
+                source_path = os.path.join(self.audio_dir, s0_wav)
+
+
+                source, fs = librosa.load(source_path, sr=self.sampling_rate, mono=True)
+                mixture, fs = librosa.load(mixture_path, sr=self.sampling_rate, mono=True)
+
+                sdr_no_sep = calculate_sdr(ref=source, est=mixture)
+                                
+                text = [labels]
+                conditions = pl_model.query_encoder.get_query_embed(
+                    modality='text',
+                    text=text,
+                    device=device 
+                )
+                    
+                input_dict = {
+                    "mixture": torch.Tensor(mixture)[None, None, :].to(device),
+                    "condition": conditions,
+                }
+                
+                sep_segment = pl_model.ss_model(input_dict)["waveform"]
+                    # sep_segment: (batch_size=1, channels_num=1, segment_samples)
+
+                sep_segment = sep_segment.squeeze(0).squeeze(0).data.cpu().numpy()
+                    # sep_segment: (segment_samples,)
+
+                sdr = calculate_sdr(ref=source, est=sep_segment)
+                sdri = sdr - sdr_no_sep
+
+                sisdr_no_sep = calculate_sisdr(ref=source, est=mixture)
+                sisdr = calculate_sisdr(ref=source, est=sep_segment)
+                sisdri = sisdr - sisdr_no_sep
+
+                sisdrs_list.append(sisdr)
+                sdris_list.append(sdri)
+                sisdris_list.append(sisdri)
+
+
+        mean_sisdr = np.mean(sisdrs_list)
+        mean_sdri = np.mean(sdris_list)
+
+        return mean_sisdr, mean_sdri

evaluation/metadata/class_labels_indices.csv

@@ -0,0 +1,528 @@
+index,mid,display_name
+0,/m/09x0r,"Speech"
+1,/m/05zppz,"Male speech, man speaking"
+2,/m/02zsn,"Female speech, woman speaking"
+3,/m/0ytgt,"Child speech, kid speaking"
+4,/m/01h8n0,"Conversation"
+5,/m/02qldy,"Narration, monologue"
+6,/m/0261r1,"Babbling"
+7,/m/0brhx,"Speech synthesizer"
+8,/m/07p6fty,"Shout"
+9,/m/07q4ntr,"Bellow"
+10,/m/07rwj3x,"Whoop"
+11,/m/07sr1lc,"Yell"
+12,/m/04gy_2,"Battle cry"
+13,/t/dd00135,"Children shouting"
+14,/m/03qc9zr,"Screaming"
+15,/m/02rtxlg,"Whispering"
+16,/m/01j3sz,"Laughter"
+17,/t/dd00001,"Baby laughter"
+18,/m/07r660_,"Giggle"
+19,/m/07s04w4,"Snicker"
+20,/m/07sq110,"Belly laugh"
+21,/m/07rgt08,"Chuckle, chortle"
+22,/m/0463cq4,"Crying, sobbing"
+23,/t/dd00002,"Baby cry, infant cry"
+24,/m/07qz6j3,"Whimper"
+25,/m/07qw_06,"Wail, moan"
+26,/m/07plz5l,"Sigh"
+27,/m/015lz1,"Singing"
+28,/m/0l14jd,"Choir"
+29,/m/01swy6,"Yodeling"
+30,/m/02bk07,"Chant"
+31,/m/01c194,"Mantra"
+32,/t/dd00003,"Male singing"
+33,/t/dd00004,"Female singing"
+34,/t/dd00005,"Child singing"
+35,/t/dd00006,"Synthetic singing"
+36,/m/06bxc,"Rapping"
+37,/m/02fxyj,"Humming"
+38,/m/07s2xch,"Groan"
+39,/m/07r4k75,"Grunt"
+40,/m/01w250,"Whistling"
+41,/m/0lyf6,"Breathing"
+42,/m/07mzm6,"Wheeze"
+43,/m/01d3sd,"Snoring"
+44,/m/07s0dtb,"Gasp"
+45,/m/07pyy8b,"Pant"
+46,/m/07q0yl5,"Snort"
+47,/m/01b_21,"Cough"
+48,/m/0dl9sf8,"Throat clearing"
+49,/m/01hsr_,"Sneeze"
+50,/m/07ppn3j,"Sniff"
+51,/m/06h7j,"Run"
+52,/m/07qv_x_,"Shuffle"
+53,/m/07pbtc8,"Walk, footsteps"
+54,/m/03cczk,"Chewing, mastication"
+55,/m/07pdhp0,"Biting"
+56,/m/0939n_,"Gargling"
+57,/m/01g90h,"Stomach rumble"
+58,/m/03q5_w,"Burping, eructation"
+59,/m/02p3nc,"Hiccup"
+60,/m/02_nn,"Fart"
+61,/m/0k65p,"Hands"
+62,/m/025_jnm,"Finger snapping"
+63,/m/0l15bq,"Clapping"
+64,/m/01jg02,"Heart sounds, heartbeat"
+65,/m/01jg1z,"Heart murmur"
+66,/m/053hz1,"Cheering"
+67,/m/028ght,"Applause"
+68,/m/07rkbfh,"Chatter"
+69,/m/03qtwd,"Crowd"
+70,/m/07qfr4h,"Hubbub, speech noise, speech babble"
+71,/t/dd00013,"Children playing"
+72,/m/0jbk,"Animal"
+73,/m/068hy,"Domestic animals, pets"
+74,/m/0bt9lr,"Dog"
+75,/m/05tny_,"Bark"
+76,/m/07r_k2n,"Yip"
+77,/m/07qf0zm,"Howl"
+78,/m/07rc7d9,"Bow-wow"
+79,/m/0ghcn6,"Growling"
+80,/t/dd00136,"Whimper (dog)"
+81,/m/01yrx,"Cat"
+82,/m/02yds9,"Purr"
+83,/m/07qrkrw,"Meow"
+84,/m/07rjwbb,"Hiss"
+85,/m/07r81j2,"Caterwaul"
+86,/m/0ch8v,"Livestock, farm animals, working animals"
+87,/m/03k3r,"Horse"
+88,/m/07rv9rh,"Clip-clop"
+89,/m/07q5rw0,"Neigh, whinny"
+90,/m/01xq0k1,"Cattle, bovinae"
+91,/m/07rpkh9,"Moo"
+92,/m/0239kh,"Cowbell"
+93,/m/068zj,"Pig"
+94,/t/dd00018,"Oink"
+95,/m/03fwl,"Goat"
+96,/m/07q0h5t,"Bleat"
+97,/m/07bgp,"Sheep"
+98,/m/025rv6n,"Fowl"
+99,/m/09b5t,"Chicken, rooster"
+100,/m/07st89h,"Cluck"
+101,/m/07qn5dc,"Crowing, cock-a-doodle-doo"
+102,/m/01rd7k,"Turkey"
+103,/m/07svc2k,"Gobble"
+104,/m/09ddx,"Duck"
+105,/m/07qdb04,"Quack"
+106,/m/0dbvp,"Goose"
+107,/m/07qwf61,"Honk"
+108,/m/01280g,"Wild animals"
+109,/m/0cdnk,"Roaring cats (lions, tigers)"
+110,/m/04cvmfc,"Roar"
+111,/m/015p6,"Bird"
+112,/m/020bb7,"Bird vocalization, bird call, bird song"
+113,/m/07pggtn,"Chirp, tweet"
+114,/m/07sx8x_,"Squawk"
+115,/m/0h0rv,"Pigeon, dove"
+116,/m/07r_25d,"Coo"
+117,/m/04s8yn,"Crow"
+118,/m/07r5c2p,"Caw"
+119,/m/09d5_,"Owl"
+120,/m/07r_80w,"Hoot"
+121,/m/05_wcq,"Bird flight, flapping wings"
+122,/m/01z5f,"Canidae, dogs, wolves"
+123,/m/06hps,"Rodents, rats, mice"
+124,/m/04rmv,"Mouse"
+125,/m/07r4gkf,"Patter"
+126,/m/03vt0,"Insect"
+127,/m/09xqv,"Cricket"
+128,/m/09f96,"Mosquito"
+129,/m/0h2mp,"Fly, housefly"
+130,/m/07pjwq1,"Buzz"
+131,/m/01h3n,"Bee, wasp, etc."
+132,/m/09ld4,"Frog"
+133,/m/07st88b,"Croak"
+134,/m/078jl,"Snake"
+135,/m/07qn4z3,"Rattle"
+136,/m/032n05,"Whale vocalization"
+137,/m/04rlf,"Music"
+138,/m/04szw,"Musical instrument"
+139,/m/0fx80y,"Plucked string instrument"
+140,/m/0342h,"Guitar"
+141,/m/02sgy,"Electric guitar"
+142,/m/018vs,"Bass guitar"
+143,/m/042v_gx,"Acoustic guitar"
+144,/m/06w87,"Steel guitar, slide guitar"
+145,/m/01glhc,"Tapping (guitar technique)"
+146,/m/07s0s5r,"Strum"
+147,/m/018j2,"Banjo"
+148,/m/0jtg0,"Sitar"
+149,/m/04rzd,"Mandolin"
+150,/m/01bns_,"Zither"
+151,/m/07xzm,"Ukulele"
+152,/m/05148p4,"Keyboard (musical)"
+153,/m/05r5c,"Piano"
+154,/m/01s0ps,"Electric piano"
+155,/m/013y1f,"Organ"
+156,/m/03xq_f,"Electronic organ"
+157,/m/03gvt,"Hammond organ"
+158,/m/0l14qv,"Synthesizer"
+159,/m/01v1d8,"Sampler"
+160,/m/03q5t,"Harpsichord"
+161,/m/0l14md,"Percussion"
+162,/m/02hnl,"Drum kit"
+163,/m/0cfdd,"Drum machine"
+164,/m/026t6,"Drum"
+165,/m/06rvn,"Snare drum"
+166,/m/03t3fj,"Rimshot"
+167,/m/02k_mr,"Drum roll"
+168,/m/0bm02,"Bass drum"
+169,/m/011k_j,"Timpani"
+170,/m/01p970,"Tabla"
+171,/m/01qbl,"Cymbal"
+172,/m/03qtq,"Hi-hat"
+173,/m/01sm1g,"Wood block"
+174,/m/07brj,"Tambourine"
+175,/m/05r5wn,"Rattle (instrument)"
+176,/m/0xzly,"Maraca"
+177,/m/0mbct,"Gong"
+178,/m/016622,"Tubular bells"
+179,/m/0j45pbj,"Mallet percussion"
+180,/m/0dwsp,"Marimba, xylophone"
+181,/m/0dwtp,"Glockenspiel"
+182,/m/0dwt5,"Vibraphone"
+183,/m/0l156b,"Steelpan"
+184,/m/05pd6,"Orchestra"
+185,/m/01kcd,"Brass instrument"
+186,/m/0319l,"French horn"
+187,/m/07gql,"Trumpet"
+188,/m/07c6l,"Trombone"
+189,/m/0l14_3,"Bowed string instrument"
+190,/m/02qmj0d,"String section"
+191,/m/07y_7,"Violin, fiddle"
+192,/m/0d8_n,"Pizzicato"
+193,/m/01xqw,"Cello"
+194,/m/02fsn,"Double bass"
+195,/m/085jw,"Wind instrument, woodwind instrument"
+196,/m/0l14j_,"Flute"
+197,/m/06ncr,"Saxophone"
+198,/m/01wy6,"Clarinet"
+199,/m/03m5k,"Harp"
+200,/m/0395lw,"Bell"
+201,/m/03w41f,"Church bell"
+202,/m/027m70_,"Jingle bell"
+203,/m/0gy1t2s,"Bicycle bell"
+204,/m/07n_g,"Tuning fork"
+205,/m/0f8s22,"Chime"
+206,/m/026fgl,"Wind chime"
+207,/m/0150b9,"Change ringing (campanology)"
+208,/m/03qjg,"Harmonica"
+209,/m/0mkg,"Accordion"
+210,/m/0192l,"Bagpipes"
+211,/m/02bxd,"Didgeridoo"
+212,/m/0l14l2,"Shofar"
+213,/m/07kc_,"Theremin"
+214,/m/0l14t7,"Singing bowl"
+215,/m/01hgjl,"Scratching (performance technique)"
+216,/m/064t9,"Pop music"
+217,/m/0glt670,"Hip hop music"
+218,/m/02cz_7,"Beatboxing"
+219,/m/06by7,"Rock music"
+220,/m/03lty,"Heavy metal"
+221,/m/05r6t,"Punk rock"
+222,/m/0dls3,"Grunge"
+223,/m/0dl5d,"Progressive rock"
+224,/m/07sbbz2,"Rock and roll"
+225,/m/05w3f,"Psychedelic rock"
+226,/m/06j6l,"Rhythm and blues"
+227,/m/0gywn,"Soul music"
+228,/m/06cqb,"Reggae"
+229,/m/01lyv,"Country"
+230,/m/015y_n,"Swing music"
+231,/m/0gg8l,"Bluegrass"
+232,/m/02x8m,"Funk"
+233,/m/02w4v,"Folk music"
+234,/m/06j64v,"Middle Eastern music"
+235,/m/03_d0,"Jazz"
+236,/m/026z9,"Disco"
+237,/m/0ggq0m,"Classical music"
+238,/m/05lls,"Opera"
+239,/m/02lkt,"Electronic music"
+240,/m/03mb9,"House music"
+241,/m/07gxw,"Techno"
+242,/m/07s72n,"Dubstep"
+243,/m/0283d,"Drum and bass"
+244,/m/0m0jc,"Electronica"
+245,/m/08cyft,"Electronic dance music"
+246,/m/0fd3y,"Ambient music"
+247,/m/07lnk,"Trance music"
+248,/m/0g293,"Music of Latin America"
+249,/m/0ln16,"Salsa music"
+250,/m/0326g,"Flamenco"
+251,/m/0155w,"Blues"
+252,/m/05fw6t,"Music for children"
+253,/m/02v2lh,"New-age music"
+254,/m/0y4f8,"Vocal music"
+255,/m/0z9c,"A capella"
+256,/m/0164x2,"Music of Africa"
+257,/m/0145m,"Afrobeat"
+258,/m/02mscn,"Christian music"
+259,/m/016cjb,"Gospel music"
+260,/m/028sqc,"Music of Asia"
+261,/m/015vgc,"Carnatic music"
+262,/m/0dq0md,"Music of Bollywood"
+263,/m/06rqw,"Ska"
+264,/m/02p0sh1,"Traditional music"
+265,/m/05rwpb,"Independent music"
+266,/m/074ft,"Song"
+267,/m/025td0t,"Background music"
+268,/m/02cjck,"Theme music"
+269,/m/03r5q_,"Jingle (music)"
+270,/m/0l14gg,"Soundtrack music"
+271,/m/07pkxdp,"Lullaby"
+272,/m/01z7dr,"Video game music"
+273,/m/0140xf,"Christmas music"
+274,/m/0ggx5q,"Dance music"
+275,/m/04wptg,"Wedding music"
+276,/t/dd00031,"Happy music"
+277,/t/dd00032,"Funny music"
+278,/t/dd00033,"Sad music"
+279,/t/dd00034,"Tender music"
+280,/t/dd00035,"Exciting music"
+281,/t/dd00036,"Angry music"
+282,/t/dd00037,"Scary music"
+283,/m/03m9d0z,"Wind"
+284,/m/09t49,"Rustling leaves"
+285,/t/dd00092,"Wind noise (microphone)"
+286,/m/0jb2l,"Thunderstorm"
+287,/m/0ngt1,"Thunder"
+288,/m/0838f,"Water"
+289,/m/06mb1,"Rain"
+290,/m/07r10fb,"Raindrop"
+291,/t/dd00038,"Rain on surface"
+292,/m/0j6m2,"Stream"
+293,/m/0j2kx,"Waterfall"
+294,/m/05kq4,"Ocean"
+295,/m/034srq,"Waves, surf"
+296,/m/06wzb,"Steam"
+297,/m/07swgks,"Gurgling"
+298,/m/02_41,"Fire"
+299,/m/07pzfmf,"Crackle"
+300,/m/07yv9,"Vehicle"
+301,/m/019jd,"Boat, Water vehicle"
+302,/m/0hsrw,"Sailboat, sailing ship"
+303,/m/056ks2,"Rowboat, canoe, kayak"
+304,/m/02rlv9,"Motorboat, speedboat"
+305,/m/06q74,"Ship"
+306,/m/012f08,"Motor vehicle (road)"
+307,/m/0k4j,"Car"
+308,/m/0912c9,"Vehicle horn, car horn, honking"
+309,/m/07qv_d5,"Toot"
+310,/m/02mfyn,"Car alarm"
+311,/m/04gxbd,"Power windows, electric windows"
+312,/m/07rknqz,"Skidding"
+313,/m/0h9mv,"Tire squeal"
+314,/t/dd00134,"Car passing by"
+315,/m/0ltv,"Race car, auto racing"
+316,/m/07r04,"Truck"
+317,/m/0gvgw0,"Air brake"
+318,/m/05x_td,"Air horn, truck horn"
+319,/m/02rhddq,"Reversing beeps"
+320,/m/03cl9h,"Ice cream truck, ice cream van"
+321,/m/01bjv,"Bus"
+322,/m/03j1ly,"Emergency vehicle"
+323,/m/04qvtq,"Police car (siren)"
+324,/m/012n7d,"Ambulance (siren)"
+325,/m/012ndj,"Fire engine, fire truck (siren)"
+326,/m/04_sv,"Motorcycle"
+327,/m/0btp2,"Traffic noise, roadway noise"
+328,/m/06d_3,"Rail transport"
+329,/m/07jdr,"Train"
+330,/m/04zmvq,"Train whistle"
+331,/m/0284vy3,"Train horn"
+332,/m/01g50p,"Railroad car, train wagon"
+333,/t/dd00048,"Train wheels squealing"
+334,/m/0195fx,"Subway, metro, underground"
+335,/m/0k5j,"Aircraft"
+336,/m/014yck,"Aircraft engine"
+337,/m/04229,"Jet engine"
+338,/m/02l6bg,"Propeller, airscrew"
+339,/m/09ct_,"Helicopter"
+340,/m/0cmf2,"Fixed-wing aircraft, airplane"
+341,/m/0199g,"Bicycle"
+342,/m/06_fw,"Skateboard"
+343,/m/02mk9,"Engine"
+344,/t/dd00065,"Light engine (high frequency)"
+345,/m/08j51y,"Dental drill, dentist's drill"
+346,/m/01yg9g,"Lawn mower"
+347,/m/01j4z9,"Chainsaw"
+348,/t/dd00066,"Medium engine (mid frequency)"
+349,/t/dd00067,"Heavy engine (low frequency)"
+350,/m/01h82_,"Engine knocking"
+351,/t/dd00130,"Engine starting"
+352,/m/07pb8fc,"Idling"
+353,/m/07q2z82,"Accelerating, revving, vroom"
+354,/m/02dgv,"Door"
+355,/m/03wwcy,"Doorbell"
+356,/m/07r67yg,"Ding-dong"
+357,/m/02y_763,"Sliding door"
+358,/m/07rjzl8,"Slam"
+359,/m/07r4wb8,"Knock"
+360,/m/07qcpgn,"Tap"
+361,/m/07q6cd_,"Squeak"
+362,/m/0642b4,"Cupboard open or close"
+363,/m/0fqfqc,"Drawer open or close"
+364,/m/04brg2,"Dishes, pots, and pans"
+365,/m/023pjk,"Cutlery, silverware"
+366,/m/07pn_8q,"Chopping (food)"
+367,/m/0dxrf,"Frying (food)"
+368,/m/0fx9l,"Microwave oven"
+369,/m/02pjr4,"Blender"
+370,/m/02jz0l,"Water tap, faucet"
+371,/m/0130jx,"Sink (filling or washing)"
+372,/m/03dnzn,"Bathtub (filling or washing)"
+373,/m/03wvsk,"Hair dryer"
+374,/m/01jt3m,"Toilet flush"
+375,/m/012xff,"Toothbrush"
+376,/m/04fgwm,"Electric toothbrush"
+377,/m/0d31p,"Vacuum cleaner"
+378,/m/01s0vc,"Zipper (clothing)"
+379,/m/03v3yw,"Keys jangling"
+380,/m/0242l,"Coin (dropping)"
+381,/m/01lsmm,"Scissors"
+382,/m/02g901,"Electric shaver, electric razor"
+383,/m/05rj2,"Shuffling cards"
+384,/m/0316dw,"Typing"
+385,/m/0c2wf,"Typewriter"
+386,/m/01m2v,"Computer keyboard"
+387,/m/081rb,"Writing"
+388,/m/07pp_mv,"Alarm"
+389,/m/07cx4,"Telephone"
+390,/m/07pp8cl,"Telephone bell ringing"
+391,/m/01hnzm,"Ringtone"
+392,/m/02c8p,"Telephone dialing, DTMF"
+393,/m/015jpf,"Dial tone"
+394,/m/01z47d,"Busy signal"
+395,/m/046dlr,"Alarm clock"
+396,/m/03kmc9,"Siren"
+397,/m/0dgbq,"Civil defense siren"
+398,/m/030rvx,"Buzzer"
+399,/m/01y3hg,"Smoke detector, smoke alarm"
+400,/m/0c3f7m,"Fire alarm"
+401,/m/04fq5q,"Foghorn"
+402,/m/0l156k,"Whistle"
+403,/m/06hck5,"Steam whistle"
+404,/t/dd00077,"Mechanisms"
+405,/m/02bm9n,"Ratchet, pawl"
+406,/m/01x3z,"Clock"
+407,/m/07qjznt,"Tick"
+408,/m/07qjznl,"Tick-tock"
+409,/m/0l7xg,"Gears"
+410,/m/05zc1,"Pulleys"
+411,/m/0llzx,"Sewing machine"
+412,/m/02x984l,"Mechanical fan"
+413,/m/025wky1,"Air conditioning"
+414,/m/024dl,"Cash register"
+415,/m/01m4t,"Printer"
+416,/m/0dv5r,"Camera"
+417,/m/07bjf,"Single-lens reflex camera"
+418,/m/07k1x,"Tools"
+419,/m/03l9g,"Hammer"
+420,/m/03p19w,"Jackhammer"
+421,/m/01b82r,"Sawing"
+422,/m/02p01q,"Filing (rasp)"
+423,/m/023vsd,"Sanding"
+424,/m/0_ksk,"Power tool"
+425,/m/01d380,"Drill"
+426,/m/014zdl,"Explosion"
+427,/m/032s66,"Gunshot, gunfire"
+428,/m/04zjc,"Machine gun"
+429,/m/02z32qm,"Fusillade"
+430,/m/0_1c,"Artillery fire"
+431,/m/073cg4,"Cap gun"
+432,/m/0g6b5,"Fireworks"
+433,/g/122z_qxw,"Firecracker"
+434,/m/07qsvvw,"Burst, pop"
+435,/m/07pxg6y,"Eruption"
+436,/m/07qqyl4,"Boom"
+437,/m/083vt,"Wood"
+438,/m/07pczhz,"Chop"
+439,/m/07pl1bw,"Splinter"
+440,/m/07qs1cx,"Crack"
+441,/m/039jq,"Glass"
+442,/m/07q7njn,"Chink, clink"
+443,/m/07rn7sz,"Shatter"
+444,/m/04k94,"Liquid"
+445,/m/07rrlb6,"Splash, splatter"
+446,/m/07p6mqd,"Slosh"
+447,/m/07qlwh6,"Squish"
+448,/m/07r5v4s,"Drip"
+449,/m/07prgkl,"Pour"
+450,/m/07pqc89,"Trickle, dribble"
+451,/t/dd00088,"Gush"
+452,/m/07p7b8y,"Fill (with liquid)"
+453,/m/07qlf79,"Spray"
+454,/m/07ptzwd,"Pump (liquid)"
+455,/m/07ptfmf,"Stir"
+456,/m/0dv3j,"Boiling"
+457,/m/0790c,"Sonar"
+458,/m/0dl83,"Arrow"
+459,/m/07rqsjt,"Whoosh, swoosh, swish"
+460,/m/07qnq_y,"Thump, thud"
+461,/m/07rrh0c,"Thunk"
+462,/m/0b_fwt,"Electronic tuner"
+463,/m/02rr_,"Effects unit"
+464,/m/07m2kt,"Chorus effect"
+465,/m/018w8,"Basketball bounce"
+466,/m/07pws3f,"Bang"
+467,/m/07ryjzk,"Slap, smack"
+468,/m/07rdhzs,"Whack, thwack"
+469,/m/07pjjrj,"Smash, crash"
+470,/m/07pc8lb,"Breaking"
+471,/m/07pqn27,"Bouncing"
+472,/m/07rbp7_,"Whip"
+473,/m/07pyf11,"Flap"
+474,/m/07qb_dv,"Scratch"
+475,/m/07qv4k0,"Scrape"
+476,/m/07pdjhy,"Rub"
+477,/m/07s8j8t,"Roll"
+478,/m/07plct2,"Crushing"
+479,/t/dd00112,"Crumpling, crinkling"
+480,/m/07qcx4z,"Tearing"
+481,/m/02fs_r,"Beep, bleep"
+482,/m/07qwdck,"Ping"
+483,/m/07phxs1,"Ding"
+484,/m/07rv4dm,"Clang"
+485,/m/07s02z0,"Squeal"
+486,/m/07qh7jl,"Creak"
+487,/m/07qwyj0,"Rustle"
+488,/m/07s34ls,"Whir"
+489,/m/07qmpdm,"Clatter"
+490,/m/07p9k1k,"Sizzle"
+491,/m/07qc9xj,"Clicking"
+492,/m/07rwm0c,"Clickety-clack"
+493,/m/07phhsh,"Rumble"
+494,/m/07qyrcz,"Plop"
+495,/m/07qfgpx,"Jingle, tinkle"
+496,/m/07rcgpl,"Hum"
+497,/m/07p78v5,"Zing"
+498,/t/dd00121,"Boing"
+499,/m/07s12q4,"Crunch"
+500,/m/028v0c,"Silence"
+501,/m/01v_m0,"Sine wave"
+502,/m/0b9m1,"Harmonic"
+503,/m/0hdsk,"Chirp tone"
+504,/m/0c1dj,"Sound effect"
+505,/m/07pt_g0,"Pulse"
+506,/t/dd00125,"Inside, small room"
+507,/t/dd00126,"Inside, large room or hall"
+508,/t/dd00127,"Inside, public space"
+509,/t/dd00128,"Outside, urban or manmade"
+510,/t/dd00129,"Outside, rural or natural"
+511,/m/01b9nn,"Reverberation"
+512,/m/01jnbd,"Echo"
+513,/m/096m7z,"Noise"
+514,/m/06_y0by,"Environmental noise"
+515,/m/07rgkc5,"Static"
+516,/m/06xkwv,"Mains hum"
+517,/m/0g12c5,"Distortion"
+518,/m/08p9q4,"Sidetone"
+519,/m/07szfh9,"Cacophony"
+520,/m/0chx_,"White noise"
+521,/m/0cj0r,"Pink noise"
+522,/m/07p_0gm,"Throbbing"
+523,/m/01jwx6,"Vibration"
+524,/m/07c52,"Television"
+525,/m/06bz3,"Radio"
+526,/m/07hvw1,"Field recording"

losses.py

@@ -0,0 +1,17 @@
+import torch
+
+
+def l1(output, target):
+    return torch.mean(torch.abs(output - target))
+
+
+def l1_wav(output_dict, target_dict):
+	return l1(output_dict['segment'], target_dict['segment'])
+
+
+def get_loss_function(loss_type):
+    if loss_type == "l1_wav":
+        return l1_wav
+
+    else:
+        raise NotImplementedError("Error!")

models/CLAP/open_clip/__init__.py

@@ -0,0 +1,25 @@
+from .factory import (
+    list_models,
+    create_model,
+    create_model_and_transforms,
+    add_model_config,
+)
+from .loss import ClipLoss, gather_features, LPLoss, lp_gather_features, LPMetrics
+from .model import (
+    CLAP,
+    CLAPTextCfg,
+    CLAPVisionCfg,
+    CLAPAudioCfp,
+    convert_weights_to_fp16,
+    trace_model,
+)
+from .openai import load_openai_model, list_openai_models
+from .pretrained import (
+    list_pretrained,
+    list_pretrained_tag_models,
+    list_pretrained_model_tags,
+    get_pretrained_url,
+    download_pretrained,
+)
+from .tokenizer import SimpleTokenizer, tokenize
+from .transform import image_transform

models/CLAP/open_clip/bert.py

@@ -0,0 +1,40 @@
+from transformers import BertTokenizer, BertModel
+
+tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+model = BertModel.from_pretrained("bert-base-uncased")
+text = "Replace me by any text you'd like."
+
+
+def bert_embeddings(text):
+    # text = "Replace me by any text you'd like."
+    encoded_input = tokenizer(text, return_tensors="pt")
+    output = model(**encoded_input)
+    return output
+
+
+from transformers import RobertaTokenizer, RobertaModel
+
+tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
+model = RobertaModel.from_pretrained("roberta-base")
+text = "Replace me by any text you'd like."
+
+
+def Roberta_embeddings(text):
+    # text = "Replace me by any text you'd like."
+    encoded_input = tokenizer(text, return_tensors="pt")
+    output = model(**encoded_input)
+    return output
+
+
+from transformers import BartTokenizer, BartModel
+
+tokenizer = BartTokenizer.from_pretrained("facebook/bart-base")
+model = BartModel.from_pretrained("facebook/bart-base")
+text = "Replace me by any text you'd like."
+
+
+def bart_embeddings(text):
+    # text = "Replace me by any text you'd like."
+    encoded_input = tokenizer(text, return_tensors="pt")
+    output = model(**encoded_input)
+    return output

models/CLAP/open_clip/bpe_simple_vocab_16e6.txt.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
+size 1356917

models/CLAP/open_clip/factory.py

@@ -0,0 +1,277 @@
+import json
+import logging
+import os
+import pathlib
+import re
+from copy import deepcopy
+from pathlib import Path
+
+import torch
+
+from .model import CLAP, convert_weights_to_fp16
+from .openai import load_openai_model
+from .pretrained import get_pretrained_url, download_pretrained
+from .transform import image_transform
+
+_MODEL_CONFIG_PATHS = [Path(__file__).parent / f"model_configs/"]
+_MODEL_CONFIGS = {}  # directory (model_name: config) of model architecture configs
+
+
+def _natural_key(string_):
+    return [int(s) if s.isdigit() else s for s in re.split(r"(\d+)", string_.lower())]
+
+
+def _rescan_model_configs():
+    global _MODEL_CONFIGS
+
+    config_ext = (".json",)
+    config_files = []
+    for config_path in _MODEL_CONFIG_PATHS:
+        if config_path.is_file() and config_path.suffix in config_ext:
+            config_files.append(config_path)
+        elif config_path.is_dir():
+            for ext in config_ext:
+                config_files.extend(config_path.glob(f"*{ext}"))
+
+    for cf in config_files:
+        if os.path.basename(cf)[0] == ".":
+            continue  # Ignore hidden files
+
+        with open(cf, "r") as f:
+            model_cfg = json.load(f)
+            if all(a in model_cfg for a in ("embed_dim", "audio_cfg", "text_cfg")):
+                _MODEL_CONFIGS[cf.stem] = model_cfg
+
+    _MODEL_CONFIGS = {
+        k: v
+        for k, v in sorted(_MODEL_CONFIGS.items(), key=lambda x: _natural_key(x[0]))
+    }
+
+
+_rescan_model_configs()  # initial populate of model config registry
+
+
+def load_state_dict(checkpoint_path: str, map_location="cpu", skip_params=True):
+    checkpoint = torch.load(checkpoint_path, map_location=map_location)
+    if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
+        state_dict = checkpoint["state_dict"]
+    else:
+        state_dict = checkpoint
+    if skip_params:
+        if next(iter(state_dict.items()))[0].startswith("module"):
+            state_dict = {k[7:]: v for k, v in state_dict.items()}
+    # for k in state_dict:
+    #     if k.startswith('transformer'):
+    #         v = state_dict.pop(k)
+    #         state_dict['text_branch.' + k[12:]] = v
+    return state_dict
+
+
+def create_model(
+    amodel_name: str,
+    tmodel_name: str,
+    pretrained: str = "",
+    precision: str = "fp32",
+    device: torch.device = torch.device("cpu"),
+    jit: bool = False,
+    force_quick_gelu: bool = False,
+    openai_model_cache_dir: str = os.path.expanduser("~/.cache/clip"),
+    skip_params=True,
+    pretrained_audio: str = "",
+    pretrained_text: str = "",
+    enable_fusion: bool = False,
+    fusion_type: str = "None"
+    # pretrained_image: bool = False,
+):
+    amodel_name = amodel_name.replace(
+        "/", "-"
+    )  # for callers using old naming with / in ViT names
+    pretrained_orig = pretrained
+    pretrained = pretrained.lower()
+    if pretrained == "openai":
+        if amodel_name in _MODEL_CONFIGS:
+            logging.info(f"Loading {amodel_name} model config.")
+            model_cfg = deepcopy(_MODEL_CONFIGS[amodel_name])
+        else:
+            logging.error(
+                f"Model config for {amodel_name} not found; available models {list_models()}."
+            )
+            raise RuntimeError(f"Model config for {amodel_name} not found.")
+
+        logging.info(f"Loading pretrained ViT-B-16 text encoder from OpenAI.")
+        # Hard Code in model name
+        model_cfg["text_cfg"]["model_type"] = tmodel_name
+        model = load_openai_model(
+            "ViT-B-16",
+            model_cfg,
+            device=device,
+            jit=jit,
+            cache_dir=openai_model_cache_dir,
+            enable_fusion=enable_fusion,
+            fusion_type=fusion_type,
+        )
+        # See https://discuss.pytorch.org/t/valueerror-attemting-to-unscale-fp16-gradients/81372
+        if precision == "amp" or precision == "fp32":
+            model = model.float()
+    else:
+        if amodel_name in _MODEL_CONFIGS:
+            logging.info(f"Loading {amodel_name} model config.")
+            model_cfg = deepcopy(_MODEL_CONFIGS[amodel_name])
+        else:
+            logging.error(
+                f"Model config for {amodel_name} not found; available models {list_models()}."
+            )
+            raise RuntimeError(f"Model config for {amodel_name} not found.")
+
+        if force_quick_gelu:
+            # override for use of QuickGELU on non-OpenAI transformer models
+            model_cfg["quick_gelu"] = True
+
+        # if pretrained_image:
+        #     if 'timm_amodel_name' in model_cfg.get('vision_cfg', {}):
+        #         # pretrained weight loading for timm models set via vision_cfg
+        #         model_cfg['vision_cfg']['timm_model_pretrained'] = True
+        #     else:
+        #         assert False, 'pretrained image towers currently only supported for timm models'
+        model_cfg["text_cfg"]["model_type"] = tmodel_name
+        model_cfg["enable_fusion"] = enable_fusion
+        model_cfg["fusion_type"] = fusion_type
+        model = CLAP(**model_cfg)
+
+        if pretrained:
+            checkpoint_path = ""
+            url = get_pretrained_url(amodel_name, pretrained)
+            if url:
+                checkpoint_path = download_pretrained(url, root=openai_model_cache_dir)
+            elif os.path.exists(pretrained_orig):
+                checkpoint_path = pretrained_orig
+            if checkpoint_path:
+                logging.info(
+                    f"Loading pretrained {amodel_name}-{tmodel_name} weights ({pretrained})."
+                )
+                ckpt = load_state_dict(checkpoint_path, skip_params=True)
+                model.load_state_dict(ckpt)
+                param_names = [n for n, p in model.named_parameters()]
+                # for n in param_names:
+                #     print(n, "\t", "Loaded" if n in ckpt else "Unloaded")
+            else:
+                logging.warning(
+                    f"Pretrained weights ({pretrained}) not found for model {amodel_name}."
+                )
+                raise RuntimeError(
+                    f"Pretrained weights ({pretrained}) not found for model {amodel_name}."
+                )
+
+        if pretrained_audio:
+            if amodel_name.startswith("PANN"):
+                if "Cnn14_mAP" in pretrained_audio:  # official checkpoint
+                    audio_ckpt = torch.load(pretrained_audio, map_location="cpu")
+                    audio_ckpt = audio_ckpt["model"]
+                    keys = list(audio_ckpt.keys())
+                    for key in keys:
+                        if (
+                            "spectrogram_extractor" not in key
+                            and "logmel_extractor" not in key
+                        ):
+                            v = audio_ckpt.pop(key)
+                            audio_ckpt["audio_branch." + key] = v
+                elif os.path.basename(pretrained_audio).startswith(
+                    "PANN"
+                ):  # checkpoint trained via HTSAT codebase
+                    audio_ckpt = torch.load(pretrained_audio, map_location="cpu")
+                    audio_ckpt = audio_ckpt["state_dict"]
+                    keys = list(audio_ckpt.keys())
+                    for key in keys:
+                        if key.startswith("sed_model"):
+                            v = audio_ckpt.pop(key)
+                            audio_ckpt["audio_branch." + key[10:]] = v
+                elif os.path.basename(pretrained_audio).startswith(
+                    "finetuned"
+                ):  # checkpoint trained via linear probe codebase
+                    audio_ckpt = torch.load(pretrained_audio, map_location="cpu")
+                else:
+                    raise ValueError("Unknown audio checkpoint")
+            elif amodel_name.startswith("HTSAT"):
+                if "HTSAT_AudioSet_Saved" in pretrained_audio:  # official checkpoint
+                    audio_ckpt = torch.load(pretrained_audio, map_location="cpu")
+                    audio_ckpt = audio_ckpt["state_dict"]
+                    keys = list(audio_ckpt.keys())
+                    for key in keys:
+                        if key.startswith("sed_model") and (
+                            "spectrogram_extractor" not in key
+                            and "logmel_extractor" not in key
+                        ):
+                            v = audio_ckpt.pop(key)
+                            audio_ckpt["audio_branch." + key[10:]] = v
+                elif os.path.basename(pretrained_audio).startswith(
+                    "HTSAT"
+                ):  # checkpoint trained via HTSAT codebase
+                    audio_ckpt = torch.load(pretrained_audio, map_location="cpu")
+                    audio_ckpt = audio_ckpt["state_dict"]
+                    keys = list(audio_ckpt.keys())
+                    for key in keys:
+                        if key.startswith("sed_model"):
+                            v = audio_ckpt.pop(key)
+                            audio_ckpt["audio_branch." + key[10:]] = v
+                elif os.path.basename(pretrained_audio).startswith(
+                    "finetuned"
+                ):  # checkpoint trained via linear probe codebase
+                    audio_ckpt = torch.load(pretrained_audio, map_location="cpu")
+                else:
+                    raise ValueError("Unknown audio checkpoint")
+            else:
+                raise f"this audio encoder pretrained checkpoint is not support"
+
+            model.load_state_dict(audio_ckpt, strict=False)
+            logging.info(
+                f"Loading pretrained {amodel_name} weights ({pretrained_audio})."
+            )
+            param_names = [n for n, p in model.named_parameters()]
+            for n in param_names:
+                print(n, "\t", "Loaded" if n in audio_ckpt else "Unloaded")
+
+        model.to(device=device)
+        if precision == "fp16":
+            assert device.type != "cpu"
+            convert_weights_to_fp16(model)
+
+        if jit:
+            model = torch.jit.script(model)
+
+    return model, model_cfg
+
+
+def create_model_and_transforms(
+    model_name: str,
+    pretrained: str = "",
+    precision: str = "fp32",
+    device: torch.device = torch.device("cpu"),
+    jit: bool = False,
+    force_quick_gelu: bool = False,
+    # pretrained_image: bool = False,
+):
+    model = create_model(
+        model_name,
+        pretrained,
+        precision,
+        device,
+        jit,
+        force_quick_gelu=force_quick_gelu,
+        # pretrained_image=pretrained_image
+    )
+    preprocess_train = image_transform(model.visual.image_size, is_train=True)
+    preprocess_val = image_transform(model.visual.image_size, is_train=False)
+    return model, preprocess_train, preprocess_val
+
+
+def list_models():
+    """enumerate available model architectures based on config files"""
+    return list(_MODEL_CONFIGS.keys())
+
+
+def add_model_config(path):
+    """add model config path or file and update registry"""
+    if not isinstance(path, Path):
+        path = Path(path)
+    _MODEL_CONFIG_PATHS.append(path)
+    _rescan_model_configs()

models/CLAP/open_clip/feature_fusion.py

@@ -0,0 +1,192 @@
+"""
+Feature Fusion for Variable-Length Data Processing
+AFF/iAFF is referred and modified from https://github.com/YimianDai/open-aff/blob/master/aff_pytorch/aff_net/fusion.py
+According to the paper: Yimian Dai et al, Attentional Feature Fusion, IEEE Winter Conference on Applications of Computer Vision, WACV 2021
+"""
+
+import torch
+import torch.nn as nn
+
+
+class DAF(nn.Module):
+    """
+    直接相加 DirectAddFuse
+    """
+
+    def __init__(self):
+        super(DAF, self).__init__()
+
+    def forward(self, x, residual):
+        return x + residual
+
+
+class iAFF(nn.Module):
+    """
+    多特征融合 iAFF
+    """
+
+    def __init__(self, channels=64, r=4, type="2D"):
+        super(iAFF, self).__init__()
+        inter_channels = int(channels // r)
+
+        if type == "1D":
+            # 本地注意力
+            self.local_att = nn.Sequential(
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+
+            # 全局注意力
+            self.global_att = nn.Sequential(
+                nn.AdaptiveAvgPool1d(1),
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+
+            # 第二次本地注意力
+            self.local_att2 = nn.Sequential(
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+            # 第二次全局注意力
+            self.global_att2 = nn.Sequential(
+                nn.AdaptiveAvgPool1d(1),
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+        elif type == "2D":
+            # 本地注意力
+            self.local_att = nn.Sequential(
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+
+            # 全局注意力
+            self.global_att = nn.Sequential(
+                nn.AdaptiveAvgPool2d(1),
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+
+            # 第二次本地注意力
+            self.local_att2 = nn.Sequential(
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+            # 第二次全局注意力
+            self.global_att2 = nn.Sequential(
+                nn.AdaptiveAvgPool2d(1),
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+        else:
+            raise f"the type is not supported"
+
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x, residual):
+        flag = False
+        xa = x + residual
+        if xa.size(0) == 1:
+            xa = torch.cat([xa, xa], dim=0)
+            flag = True
+        xl = self.local_att(xa)
+        xg = self.global_att(xa)
+        xlg = xl + xg
+        wei = self.sigmoid(xlg)
+        xi = x * wei + residual * (1 - wei)
+
+        xl2 = self.local_att2(xi)
+        xg2 = self.global_att(xi)
+        xlg2 = xl2 + xg2
+        wei2 = self.sigmoid(xlg2)
+        xo = x * wei2 + residual * (1 - wei2)
+        if flag:
+            xo = xo[0].unsqueeze(0)
+        return xo
+
+
+class AFF(nn.Module):
+    """
+    多特征融合 AFF
+    """
+
+    def __init__(self, channels=64, r=4, type="2D"):
+        super(AFF, self).__init__()
+        inter_channels = int(channels // r)
+
+        if type == "1D":
+            self.local_att = nn.Sequential(
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+            self.global_att = nn.Sequential(
+                nn.AdaptiveAvgPool1d(1),
+                nn.Conv1d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv1d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm1d(channels),
+            )
+        elif type == "2D":
+            self.local_att = nn.Sequential(
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+            self.global_att = nn.Sequential(
+                nn.AdaptiveAvgPool2d(1),
+                nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(inter_channels),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+                nn.BatchNorm2d(channels),
+            )
+        else:
+            raise f"the type is not supported."
+
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x, residual):
+        flag = False
+        xa = x + residual
+        if xa.size(0) == 1:
+            xa = torch.cat([xa, xa], dim=0)
+            flag = True
+        xl = self.local_att(xa)
+        xg = self.global_att(xa)
+        xlg = xl + xg
+        wei = self.sigmoid(xlg)
+        xo = 2 * x * wei + 2 * residual * (1 - wei)
+        if flag:
+            xo = xo[0].unsqueeze(0)
+        return xo