Upload 544 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+example.wav filter=lfs diff=lfs merge=lfs -text
+example2.wav filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Jiarui Hai
+
+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.

README.md

@@ -1,14 +1,60 @@
----
-title: FlexSED
-emoji: 🏆
-colorFrom: purple
-colorTo: purple
-sdk: gradio
-sdk_version: 5.49.1
-app_file: app.py
-pinned: false
-license: mit
-short_description: 'FlexSED: An Open-Vocabulary Sound Event Detection System'
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# FlexSED: Towards Open-Vocabulary Sound Event Detection
+
+[![arXiv](https://img.shields.io/badge/arXiv-2409.10819-brightgreen.svg?style=flat-square)](https://arxiv.org/abs/2509.18606)
+[![Hugging Face Models](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Models-blue)](https://huggingface.co/Higobeatz/FlexSED/tree/main)
+
+
+## News
+- Oct 2025: 📦 Released code and pretrained checkpoint  
+- Sep 2025: 🎉 FlexSED Spotlighted at WASPAA 2025
+
+
+## Installation
+
+Clone the repository:
+```
+git clone [email protected]:JHU-LCAP/FlexSED.git
+```
+Install the dependencies:
+```
+cd FlexSED
+pip install -r requirements.txt
+```
+
+## Usage
+```python
+from api import FlexSED
+import torch
+import soundfile as sf
+
+# load model
+flexsed = FlexSED(device='cuda')
+
+# run inference
+events = ["Dog"]
+preds = flexsed.run_inference("example.wav", events)
+
+# visualize prediciton
+flexsed.to_multi_plot(preds, events, fname="example2")
+
+# (Optional) visualize prediciton by video
+# flexsed.to_multi_video(preds, events, audio_path="example2.wav", fname="example2")
+```
+
+## Training
+
+WIP
+
+
+## Reference
+
+If you find the code useful for your research, please consider citing:
+
+```bibtex
+@article{hai2025flexsed,
+  title={FlexSED: Towards Open-Vocabulary Sound Event Detection},
+  author={Hai, Jiarui and Wang, Helin and Guo, Weizhe and Elhilali, Mounya},
+  journal={arXiv preprint arXiv:2509.18606},
+  year={2025}
+}
+```

api.py

@@ -0,0 +1,185 @@
+import torch
+import librosa
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+from transformers import AutoTokenizer, ClapTextModelWithProjection
+from src.models.transformer import Dasheng_Encoder
+from src.models.sed_decoder import Decoder, TSED_Wrapper
+from src.utils import load_yaml_with_includes
+
+
+class FlexSED:
+    def __init__(
+        self,
+        config_path='src/configs/model.yml',
+        ckpt_path='ckpts/flexsed_as.pt',
+        ckpt_url='https://huggingface.co/Higobeatz/FlexSED/resolve/main/ckpts/flexsed_as.pt',
+        device='cuda'
+    ):
+        """
+        Initialize FlexSED with model, CLAP, and tokenizer loaded once.
+        If the checkpoint is not available locally, it will be downloaded automatically.
+        """
+        self.device = device
+        params = load_yaml_with_includes(config_path)
+
+        # Ensure checkpoint exists
+        if not os.path.exists(ckpt_path):
+            print(f"[FlexSED] Downloading checkpoint from {ckpt_url} ...")
+            state_dict = torch.hub.load_state_dict_from_url(ckpt_url, map_location="cpu")
+        else:
+            state_dict = torch.load(ckpt_path, map_location="cpu")
+
+        # Encoder + Decoder
+        encoder = Dasheng_Encoder(**params['encoder']).to(self.device)
+        decoder = Decoder(**params['decoder']).to(self.device)
+        self.model = TSED_Wrapper(encoder, decoder, params['ft_blocks'], params['frozen_encoder'])
+        self.model.load_state_dict(state_dict['model'])
+        self.model.eval()
+
+        # CLAP text model
+        self.clap = ClapTextModelWithProjection.from_pretrained("laion/clap-htsat-unfused")
+        self.clap.eval()
+        self.tokenizer = AutoTokenizer.from_pretrained("laion/clap-htsat-unfused")
+
+    def run_inference(self, audio_path, events, norm_audio=True):
+        """
+        Run inference on audio for given events.
+        """
+        audio, sr = librosa.load(audio_path, sr=16000)
+        audio = torch.tensor([audio]).to(self.device)
+
+        if norm_audio:
+            eps = 1e-9
+            max_val = torch.max(torch.abs(audio))
+            audio = audio / (max_val + eps)
+
+        clap_embeds = []
+        with torch.no_grad():
+            for event in events:
+                text = f"The sound of {event.replace('_',' ')}"
+                inputs = self.tokenizer([text], padding=True, return_tensors="pt")
+                outputs = self.clap(**inputs)
+                text_embeds = outputs.text_embeds.unsqueeze(1)
+                clap_embeds.append(text_embeds)
+
+            query = torch.cat(clap_embeds, dim=1).to(self.device)
+            mel = self.model.forward_to_spec(audio)
+            preds = self.model(mel, query)
+            preds = torch.sigmoid(preds).cpu()
+
+        return preds  # shape: [num_events, 1, T]
+
+    # ---------- Multi-event plotting ----------
+    @staticmethod
+    def plot_and_save_multi(preds, events, sr=25, out_dir="./plots", fname="all_events"):
+        os.makedirs(out_dir, exist_ok=True)
+        preds_np = preds.squeeze(1).numpy()  # [num_events, T]
+        T = preds_np.shape[1]
+
+        plt.figure(figsize=(12, len(events) * 0.6 + 2))
+        plt.imshow(
+            preds_np,
+            aspect="auto",
+            cmap="Blues",
+            extent=[0, T/sr, 0, len(events)],
+            vmin=0, vmax=1, origin="lower"
+
+        )
+        plt.colorbar(label="Probability")
+        plt.yticks(np.arange(len(events)) + 0.5, events)
+        plt.xlabel("Time (s)")
+        plt.ylabel("Events")
+        plt.title("Event Predictions")
+
+        save_path = os.path.join(out_dir, f"{fname}.png")
+        plt.savefig(save_path, dpi=200, bbox_inches="tight")
+        plt.close()
+        return save_path
+
+    def to_multi_plot(self, preds, events, out_dir="./plots", fname="all_events"):
+        return self.plot_and_save_multi(preds, events, out_dir=out_dir, fname=fname)
+
+    # ---------- Multi-event video ----------
+    @staticmethod
+    def make_multi_event_video(preds, events, sr=25, out_dir="./videos",
+                               audio_path=None, fps=25, highlight=True, fname="all_events"):
+        from moviepy.editor import ImageSequenceClip, AudioFileClip
+        from tqdm import tqdm
+
+        os.makedirs(out_dir, exist_ok=True)
+        preds_np = preds.squeeze(1).numpy()  # [num_events, T]
+        T = preds_np.shape[1]
+        duration = T / sr
+
+        frames = []
+        n_frames = int(duration * fps)
+
+        for i in tqdm(range(n_frames)):
+            t = int(i * T / n_frames)
+            plt.figure(figsize=(12, len(events) * 0.6 + 2))
+
+            if highlight:
+                mask = np.zeros_like(preds_np)
+                mask[:, :t+1] = preds_np[:, :t+1]
+                plt.imshow(
+                    mask,
+                    aspect="auto",
+                    cmap="Blues",
+                    extent=[0, T/sr, 0, len(events)],
+                    vmin=0, vmax=1, origin="lower"
+                )
+            else:
+                plt.imshow(
+                    preds_np[:, :t+1],
+                    aspect="auto",
+                    cmap="Blues",
+                    extent=[0, (t+1)/sr, 0, len(events)],
+                    vmin=0, vmax=1, origin="lower"
+                )
+
+            plt.colorbar(label="Probability")
+            plt.yticks(np.arange(len(events)) + 0.5, events)
+            plt.xlabel("Time (s)")
+            plt.ylabel("Events")
+            plt.title("Event Predictions")
+
+            frame_path = f"/tmp/frame_{i:04d}.png"
+            plt.savefig(frame_path, dpi=150, bbox_inches="tight")
+            plt.close()
+            frames.append(frame_path)
+
+        clip = ImageSequenceClip(frames, fps=fps)
+        if audio_path is not None:
+            audio = AudioFileClip(audio_path).subclip(0, duration)
+            clip = clip.set_audio(audio)
+
+        save_path = os.path.join(out_dir, f"{fname}.mp4")
+        clip.write_videofile(
+            save_path,
+            fps=fps,
+            codec="mpeg4",
+            audio_codec="aac"
+        )
+
+        for f in frames:
+            os.remove(f)
+
+        return save_path
+
+    def to_multi_video(self, preds, events, audio_path, out_dir="./videos", fname="all_events"):
+        return self.make_multi_event_video(
+            preds, events, audio_path=audio_path, out_dir=out_dir, fname=fname
+        )
+
+
+if __name__ == "__main__":
+    flexsed = FlexSED(device='cuda')
+
+    events = ["Door", "Laughter", "Dog"]
+    preds = flexsed.run_inference("example2.wav", events)
+
+    # Combined plot & video
+    flexsed.to_multi_plot(preds, events, fname="example2")
+    # flexsed.to_multi_video(preds, events, audio_path="example2.wav", fname="example2")

app.py

@@ -1,7 +1,74 @@
 import gradio as gr
+import torch
+from api import FlexSED
+import tempfile
+import os
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Load model once on startup
+flexsed = FlexSED(device="cuda" if torch.cuda.is_available() else "cpu")
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+def run_flexsed(audio_file, event_list):
+    """
+    Run inference using FlexSED and return prediction plot.
+    """
+    if not audio_file:
+        return None
+
+    # Split events by semicolon or comma
+    events = [e.strip() for e in event_list.split(";") if e.strip()]
+    if not events:
+        return None
+
+    # Run inference
+    preds = flexsed.run_inference(audio_file, events)
+
+    # Generate visualization
+    output_fname = os.path.join(tempfile.gettempdir(), "flexsed_output")
+    flexsed.to_multi_plot(preds, events, fname=output_fname)
+    plot_path = f"{output_fname}.png"
+
+    return plot_path
+
+
+# App layout
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as app:
+    # Header
+    gr.Markdown("""
+    ## 🎧 FlexSED: A Flexible Open-Vocabulary Sound Event Detection System
+
+    👋 Welcome to the **FlexSED live demo** — explore **prompt-guided sound event detection** in real audio clips.
+
+    🔗 Learn more on the [FlexSED GitHub Repository](https://github.com/JHU-LCAP/FlexSED)
+    """)
+
+    gr.Markdown("### 🔍 Upload or choose an example below to detect sound events:")
+
+    with gr.Row():
+        # Left column: Inputs
+        with gr.Column(scale=1):
+            audio_input = gr.Audio(type="filepath", label="🎵 Upload Audio (.wav)")
+            text_input = gr.Textbox(label="Event list (semicolon-separated)", value="Male speech; Door; Dog; Laughter")
+
+            with gr.Row():
+                detect_btn = gr.Button("🎯 Detect", variant="primary")
+                clear_btn = gr.Button("🧹 Clear")
+
+        # Right column: Output
+        with gr.Column(scale=1):
+            image_output = gr.Image(label="Prediction Plot", show_label=True, elem_id="output-image")
+            gr.Examples(
+                examples=[
+                    ["example.wav", "Male speech; Door; Dog; Laughter"],
+                    ["example2.wav", "Male speech; Bee; Gunshot, gunfire"],
+                ],
+                inputs=[audio_input, text_input],
+                label="Example Audios"
+            )
+
+    # Function bindings
+    detect_btn.click(run_flexsed, inputs=[audio_input, text_input], outputs=image_output)
+    clear_btn.click(lambda: (None, "Male speech; Door; Dog; Laughter"), outputs=[audio_input, text_input])
+
+
+if __name__ == "__main__":
+    app.launch(share=True)

example.wav

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

example2.wav

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

requirements.txt

@@ -0,0 +1,6 @@
+gradio
+torch
+soundfile
+matplotlib
+numpy
+librosa

src/.ipynb_checkpoints/prepare_clap-checkpoint.py

@@ -0,0 +1,39 @@
+import os
+import pandas as pd
+import torch
+from transformers import AutoTokenizer, ClapTextModelWithProjection
+
+if __name__ == '__main__':
+    # Load the CLAP model and tokenizer
+    model = ClapTextModelWithProjection.from_pretrained("laion/clap-htsat-unfused")
+    model.eval()
+    tokenizer = AutoTokenizer.from_pretrained("laion/clap-htsat-unfused")
+
+    # Path to the input CSV file
+    input_csv_path = '/home/user/SSD/Dataset/Audioset_SL/no_rule_all/label_to_id.csv'
+    output_path = 'clap_embedding/'  # Replace with your desired output folder path
+
+    # Create the output folder if it doesn't exist
+    os.makedirs(output_path, exist_ok=True)
+
+    # Read the CSV file
+    df = pd.read_csv(input_csv_path)
+
+    # Get unique event labels
+    events = df['label'].unique()
+
+    with torch.no_grad():  # Disable gradient computation
+        # Process each event
+        for event in events:
+            text = event.replace('_', ' ')  # Replace underscores with spaces
+            text = f'The sound of {text}'
+            print(text)
+            inputs = tokenizer([text], padding=True, return_tensors="pt")
+            outputs = model(**inputs)
+            text_embeds = outputs.text_embeds
+
+            # Save the embeddings to a .pt file
+            output_file = os.path.join(output_path, f"{event}.pt")
+            torch.save(text_embeds, output_file)
+
+        print("Embedding extraction and saving complete!")

src/.ipynb_checkpoints/test-checkpoint.py

@@ -0,0 +1,140 @@
+import random
+import argparse
+import os
+import time
+import numpy as np
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator
+
+from models.transformer import Dasheng_Encoder
+from models.sed_decoder import Decoder, TSED_Wrapper
+from dataset.tsed import TSED_AS
+from dataset.tsed_val import TSED_Val
+from utils import load_yaml_with_includes, get_lr_scheduler, ConcatDatasetBatchSampler
+from utils.data_aug import frame_shift, mixup, time_mask, feature_transformation
+from val import val_psds
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+
+    # Config settings
+    parser.add_argument('--config-name', type=str, default='configs/model.yml')
+    parser.add_argument('--ckpt', type=str, default='20000.pt')
+
+    # Training settings
+    parser.add_argument("--amp", type=str, default='fp16')
+    parser.add_argument('--epochs', type=int, default=20)
+    parser.add_argument('--num-workers', type=int, default=8)
+    parser.add_argument('--num-threads', type=int, default=1)
+    parser.add_argument('--eval-every-step', type=int, default=5000)
+    parser.add_argument('--save-every-step', type=int, default=5000)
+    # parser.add_argument('--dataloader', type=str, default='EACaps')
+    parser.add_argument("--logit-normal-indices", type=bool, default=False)
+
+    # Log and random seed
+    parser.add_argument('--random-seed', type=int, default=2024)
+    parser.add_argument('--log-step', type=int, default=100)
+    parser.add_argument('--log-dir', type=str, default='../logs/')
+    parser.add_argument('--save-dir', type=str, default='../ckpts/')
+    return parser.parse_args()
+
+
+def setup_directories(args, params):
+    args.log_dir = os.path.join(args.log_dir, params['model_name']) + '/'
+    args.save_dir = os.path.join(args.save_dir, params['model_name']) + '/'
+
+    os.makedirs(args.log_dir, exist_ok=True)
+    os.makedirs(args.save_dir, exist_ok=True)
+
+
+def set_device(args):
+    torch.set_num_threads(args.num_threads)
+    if torch.cuda.is_available():
+        args.device = 'cuda'
+        torch.cuda.manual_seed_all(args.random_seed)
+        torch.backends.cuda.matmul.allow_tf32 = True
+        if torch.backends.cudnn.is_available():
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+    else:
+        args.device = 'cpu'
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    params = load_yaml_with_includes(args.config_name)
+    set_device(args)
+    setup_directories(args, params)
+
+    random.seed(args.random_seed)
+    torch.manual_seed(args.random_seed)
+
+    # use accelerator for multi-gpu training
+    accelerator = Accelerator(mixed_precision=args.amp,
+                              gradient_accumulation_steps=params['opt']['accumulation_steps'],
+                              step_scheduler_with_optimizer=False)
+
+    train_set = TSED_AS(**params['data']['train_data'])
+    train_loader = DataLoader(train_set, batch_size=params['opt']['batch_size'], num_workers=args.num_workers)
+
+    # val_set = TSED_Val(**params['data']['val_data'])
+    # val_loader = DataLoader(val_set, num_workers=0, batch_size=1, shuffle=False)
+
+    test_set = TSED_Val(**params['data']['test_data'])
+    test_loader = DataLoader(test_set, num_workers=0, batch_size=1, shuffle=False)
+
+    encoder = Dasheng_Encoder(**params['encoder']).to(accelerator.device)
+    pretrained_url = 'https://zenodo.org/records/11511780/files/dasheng_base.pt?download=1'
+    dump = torch.hub.load_state_dict_from_url(pretrained_url, map_location='cpu')
+    model_parmeters = dump['model']
+    # pretrained_url = 'https://zenodo.org/records/13315686/files/dasheng_audioset_mAP497.pt?download=1'
+    # dump = torch.hub.load_state_dict_from_url(pretrained_url, map_location='cpu')
+    # model_parmeters = dump
+    encoder.load_state_dict(model_parmeters)
+
+    decoder = Decoder(**params['decoder']).to(accelerator.device)
+
+    model = TSED_Wrapper(encoder, decoder, params['ft_blocks'], params['frozen_encoder'])
+    print(f"Trainable Parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.2f}M")
+
+    model.load_state_dict(torch.load(args.ckpt, map_location='cpu')['model'])
+
+    if params['frozen_encoder']:
+        optimizer = torch.optim.AdamW(
+            model.parameters(),
+            lr=params['opt']['learning_rate'],
+            weight_decay=params['opt']['weight_decay'],
+            betas=(params['opt']['beta1'], params['opt']['beta2']),
+            eps=params['opt']['adam_epsilon'])
+    else:
+        optimizer = torch.optim.AdamW(
+            [
+                {'params': model.encoder.parameters(), 'lr': 0.1 * params['opt']['learning_rate']},
+                {'params': model.decoder.parameters(), 'lr': params['opt']['learning_rate']}
+            ],
+            weight_decay=params['opt']['weight_decay'],
+            betas=(params['opt']['beta1'], params['opt']['beta2']),
+            eps=params['opt']['adam_epsilon'])
+
+    lr_scheduler = get_lr_scheduler(optimizer, 'customized', **params['opt']['lr_scheduler'])
+
+    strong_loss_func = nn.BCEWithLogitsLoss()
+
+    model, optimizer, lr_scheduler, train_loader, test_loader = accelerator.prepare(
+        model, optimizer, lr_scheduler, train_loader, test_loader)
+
+    global_step = 0.0
+    losses = 0.0
+
+    if accelerator.is_main_process:
+        model_module = model.module if hasattr(model, 'module') else model
+        val_psds(model_module, test_loader, params, epoch='test_full', split='test',
+                 save_path=args.log_dir + 'output/', device=accelerator.device)

src/.ipynb_checkpoints/train-checkpoint.py

@@ -0,0 +1,208 @@
+import random
+import argparse
+import os
+import time
+import numpy as np
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator
+
+from models.transformer import Dasheng_Encoder
+from models.sed_decoder import Decoder, TSED_Wrapper
+from dataset.tsed import TSED_AS
+from dataset.tsed_val import TSED_Val
+from utils import load_yaml_with_includes, get_lr_scheduler, ConcatDatasetBatchSampler
+from utils.data_aug import frame_shift, mixup, time_mask, feature_transformation
+from val import val_psds
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+
+    # Config settings
+    parser.add_argument('--config-name', type=str, default='configs/model.yml')
+
+    # Training settings
+    parser.add_argument("--amp", type=str, default='fp16')
+    parser.add_argument('--epochs', type=int, default=20)
+    parser.add_argument('--num-workers', type=int, default=8)
+    parser.add_argument('--num-threads', type=int, default=1)
+    parser.add_argument('--eval-every-step', type=int, default=5000)
+    parser.add_argument('--save-every-step', type=int, default=5000)
+    # parser.add_argument('--dataloader', type=str, default='EACaps')
+    parser.add_argument("--logit-normal-indices", type=bool, default=False)
+
+    # Log and random seed
+    parser.add_argument('--random-seed', type=int, default=2024)
+    parser.add_argument('--log-step', type=int, default=100)
+    parser.add_argument('--log-dir', type=str, default='../logs/')
+    parser.add_argument('--save-dir', type=str, default='../ckpts/')
+    return parser.parse_args()
+
+
+def setup_directories(args, params):
+    args.log_dir = os.path.join(args.log_dir, params['model_name']) + '/'
+    args.save_dir = os.path.join(args.save_dir, params['model_name']) + '/'
+
+    os.makedirs(args.log_dir, exist_ok=True)
+    os.makedirs(args.save_dir, exist_ok=True)
+
+
+def set_device(args):
+    torch.set_num_threads(args.num_threads)
+    if torch.cuda.is_available():
+        args.device = 'cuda'
+        torch.cuda.manual_seed_all(args.random_seed)
+        torch.backends.cuda.matmul.allow_tf32 = True
+        if torch.backends.cudnn.is_available():
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+    else:
+        args.device = 'cpu'
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    params = load_yaml_with_includes(args.config_name)
+    set_device(args)
+    setup_directories(args, params)
+
+    random.seed(args.random_seed)
+    torch.manual_seed(args.random_seed)
+
+    # use accelerator for multi-gpu training
+    accelerator = Accelerator(mixed_precision=args.amp,
+                              gradient_accumulation_steps=params['opt']['accumulation_steps'],
+                              step_scheduler_with_optimizer=False)
+
+    train_set = TSED_AS(**params['data']['train_data'])
+    train_loader = DataLoader(train_set, shuffle=True,
+                              batch_size=params['opt']['batch_size'],
+                              num_workers=args.num_workers)
+
+    val_set = TSED_Val(**params['data']['val_data'])
+    val_loader = DataLoader(val_set, num_workers=0, batch_size=1, shuffle=False)
+
+    # test_set = TSED_Val(**params['data']['test_data'])
+    # test_loader = DataLoader(val_set, num_workers=0, batch_size=1, shuffle=False)
+
+    encoder = Dasheng_Encoder(**params['encoder']).to(accelerator.device)
+    pretrained_url = 'https://zenodo.org/records/11511780/files/dasheng_base.pt?download=1'
+    dump = torch.hub.load_state_dict_from_url(pretrained_url, map_location='cpu')
+    model_parmeters = dump['model']
+    # pretrained_url = 'https://zenodo.org/records/13315686/files/dasheng_audioset_mAP497.pt?download=1'
+    # dump = torch.hub.load_state_dict_from_url(pretrained_url, map_location='cpu')
+    # model_parmeters = dump
+    encoder.load_state_dict(model_parmeters)
+
+    decoder = Decoder(**params['decoder']).to(accelerator.device)
+
+    model = TSED_Wrapper(encoder, decoder, params['ft_blocks'], params['frozen_encoder'])
+    print(f"Trainable Parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.2f}M")
+
+    # model.load_state_dict(torch.load('../ckpts/TSED_AS_filter/20000.0.pt', map_location='cpu')['model'])
+
+    if params['frozen_encoder']:
+        optimizer = torch.optim.AdamW(
+            model.parameters(),
+            lr=params['opt']['learning_rate'],
+            weight_decay=params['opt']['weight_decay'],
+            betas=(params['opt']['beta1'], params['opt']['beta2']),
+            eps=params['opt']['adam_epsilon'])
+    else:
+        optimizer = torch.optim.AdamW(
+            [
+                {'params': model.encoder.parameters(), 'lr': 0.1 * params['opt']['learning_rate']},
+                {'params': model.decoder.parameters(), 'lr': params['opt']['learning_rate']}
+            ],
+            weight_decay=params['opt']['weight_decay'],
+            betas=(params['opt']['beta1'], params['opt']['beta2']),
+            eps=params['opt']['adam_epsilon'])
+
+    lr_scheduler = get_lr_scheduler(optimizer, 'customized', **params['opt']['lr_scheduler'])
+
+    strong_loss_func = nn.BCEWithLogitsLoss()
+
+    model, optimizer, lr_scheduler, train_loader, val_loader = accelerator.prepare(
+        model, optimizer, lr_scheduler, train_loader, val_loader)
+
+    global_step = 0.0
+    losses = 0.0
+
+    if accelerator.is_main_process:
+        model_module = model.module if hasattr(model, 'module') else model
+        val_psds(model_module, val_loader, params, epoch='debug', split='val',
+                 save_path=args.log_dir + 'output/', device=accelerator.device)
+
+    for epoch in range(args.epochs):
+        model.train()
+        for step, batch in enumerate(tqdm(train_loader)):
+            with accelerator.accumulate(model):
+                audio, cls, label, _ = batch
+                mel = model.forward_to_spec(audio)
+
+                # data aug
+                mel, label = frame_shift(mel, label, params['net_pooling'])
+                mel, label = time_mask(mel, label, params["net_pooling"],
+                                       mask_ratios=params['data_aug']["time_mask_ratios"])
+                mel, _ = feature_transformation(mel, **params['data_aug']["transform"])
+
+                strong_pred = model(mel, cls)
+
+                B, N, L = label.shape
+                label = label.reshape(B * N, L)
+                label = label.unsqueeze(1)
+
+                loss = strong_loss_func(strong_pred, label)
+
+                accelerator.backward(loss)
+
+                # clip grad up
+                if accelerator.sync_gradients:
+                    if 'grad_clip' in params['opt'] and params['opt']['grad_clip'] > 0:
+                        accelerator.clip_grad_norm_(model.parameters(),
+                                                    max_norm=params['opt']['grad_clip'])
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+                global_step += 1/params['opt']['accumulation_steps']
+                losses += loss.item()/params['opt']['accumulation_steps']
+
+            if accelerator.is_main_process:
+                if global_step % args.log_step == 0:
+                    current_time = time.asctime(time.localtime(time.time()))
+                    epoch_info = f'Epoch: [{epoch + 1}][{args.epochs}]'
+                    batch_info = f'Global Step: {global_step}'
+                    loss_info = f'Loss: {losses / args.log_step:.6f}'
+
+                    # Extract the learning rate from the optimizer
+                    lr = optimizer.param_groups[0]['lr']
+                    lr_info = f'Learning Rate: {lr:.6f}'
+
+                    log_message = f'{current_time}\n{epoch_info}    {batch_info}    {loss_info}    {lr_info}\n'
+
+                    with open(args.log_dir + 'log.txt', mode='a') as n:
+                        n.write(log_message)
+
+                    losses = 0.0
+
+            # check performance
+            if (global_step + 1) % args.eval_every_step == 0:
+                if accelerator.is_main_process:
+                    model_module = model.module if hasattr(model, 'module') else model
+                    val_psds(model_module, val_loader, params, epoch=global_step+1, split='val',
+                             save_path=args.log_dir + 'output/', device=accelerator.device)
+                    # save model
+                    unwrapped_model = accelerator.unwrap_model(model)
+                    accelerator.save({
+                        "model": model.state_dict(),
+                    }, args.save_dir + str(global_step+1) + '.pt')
+                accelerator.wait_for_everyone()
+                model.train()

src/.ipynb_checkpoints/val-checkpoint.py

@@ -0,0 +1,141 @@
+import torch
+import os
+import pandas as pd
+from tqdm import tqdm
+import sed_scores_eval
+from desed_task.evaluation.evaluation_measures import (compute_per_intersection_macro_f1,
+                                                       compute_psds_from_operating_points,
+                                                       compute_psds_from_scores)
+from local.utils import (batched_decode_preds,)
+from utils.sed import Encoder
+import numpy as np
+
+
+@torch.no_grad()
+def val_psds(model, val_loader, params, epoch, split, save_path, device):
+    label_df = pd.read_csv(params['data'][split]['label'])
+    EVENTS = label_df['label'].tolist()
+
+    clap_emb = []
+    for event in EVENTS:
+        cls = torch.load(params['data']['train_data']['clap_dir'] + event + '.pt').to(device)
+        cls = cls.unsqueeze(1)
+        clap_emb.append(cls)
+    cls = torch.cat(clap_emb, dim=1)
+
+    encoder = Encoder(EVENTS, audio_len=10, frame_len=160, frame_hop=160, net_pooling=4, sr=16000)
+
+    model.eval()
+    test_csv = params['data'][split]["csv"]
+    test_dur = params['data'][split]["dur"]
+
+    gt = pd.read_csv(test_csv, sep='\t')
+
+    test_scores_postprocessed_buffer = {}
+    test_scores_postprocessed_buffer_tsed = {}
+    test_thresholds = [0.5]
+    test_psds_buffer = {k: pd.DataFrame() for k in test_thresholds}
+    test_psds_buffer_tsed = {k: pd.DataFrame() for k in test_thresholds}
+
+    for batch in tqdm(val_loader):
+        audio, filenames = batch
+        B = audio.shape[0]
+        N = cls.shape[1]
+        cls = cls.expand(B, -1, -1)
+
+        audio = audio.to(device)
+        mel = model.forward_to_spec(audio)
+
+        preds = model(mel, cls)
+        preds = torch.sigmoid(preds)
+        preds = preds.reshape(B, N, -1)
+        preds_tsed = preds.clone()
+        # tsed assumes sound exitencance is known
+        for idx, filename in enumerate(filenames):
+            weak_label = list(gt[gt['filename'] == filename]['event_label'].unique())
+            for j, event in enumerate(EVENTS):
+                if event not in weak_label:
+                    preds_tsed[idx][j] = 0.0
+        # preds = preds.transpose(1, 2)
+
+        (_, scores_postprocessed_strong, _,) = \
+            batched_decode_preds(
+                preds,
+                filenames,
+                encoder,
+                median_filter=9,
+                thresholds=list(test_psds_buffer.keys()), )
+        test_scores_postprocessed_buffer.update(scores_postprocessed_strong)
+
+        (_, scores_postprocessed_strong_tsed, _,) = \
+            batched_decode_preds(
+                preds_tsed,
+                filenames,
+                encoder,
+                median_filter=9,
+                thresholds=list(test_psds_buffer_tsed.keys()), )
+        test_scores_postprocessed_buffer_tsed.update(scores_postprocessed_strong_tsed)
+
+    ground_truth = sed_scores_eval.io.read_ground_truth_events(test_csv)
+    audio_durations = sed_scores_eval.io.read_audio_durations(test_dur)
+
+    ground_truth = {
+        audio_id: ground_truth[audio_id]
+        for audio_id in test_scores_postprocessed_buffer
+    }
+    audio_durations = {
+        audio_id: audio_durations[audio_id]
+        for audio_id in test_scores_postprocessed_buffer
+    }
+
+    psds1_sed_scores_eval, psds1_cls = compute_psds_from_scores(
+        test_scores_postprocessed_buffer,
+        ground_truth,
+        audio_durations,
+        dtc_threshold=0.7,
+        gtc_threshold=0.7,
+        cttc_threshold=None,
+        alpha_ct=0.0,
+        alpha_st=0.0,
+        # save_dir=os.path.join(save_dir, "student", "scenario1"),
+    )
+    psds1_cls['overall'] = psds1_sed_scores_eval
+    psds1_cls['macro_averaged'] = np.array([v for k, v in psds1_cls.items()]).mean()
+    psds1_cls['name'] = 'psds1'
+
+    psds1_sed_scores_eval_tsed, psds1_cls_tsed = compute_psds_from_scores(
+        test_scores_postprocessed_buffer_tsed,
+        ground_truth,
+        audio_durations,
+        dtc_threshold=0.7,
+        gtc_threshold=0.7,
+        cttc_threshold=None,
+        alpha_ct=0.0,
+        alpha_st=0.0,
+        # save_dir=os.path.join(save_dir, "student", "scenario1"),
+    )
+
+    psds1_cls_tsed['overall'] = psds1_sed_scores_eval_tsed
+    psds1_cls_tsed['macro_averaged'] = np.array([v for k, v in psds1_cls_tsed.items()]).mean()
+    psds1_cls_tsed['name'] = 'psds1_tsed'
+
+    # psds2_sed_scores_eval, psds2_cls = compute_psds_from_scores(
+    #     test_scores_postprocessed_buffer,
+    #     ground_truth,
+    #     audio_durations,
+    #     dtc_threshold=0.1,
+    #     gtc_threshold=0.1,
+    #     cttc_threshold=0.3,
+    #     alpha_ct=0.5,
+    #     alpha_st=1,
+    #     # save_dir=os.path.join(save_dir, "student", "scenario1"),
+    # )
+    # psds2_cls['overall'] = psds2_sed_scores_eval
+    # psds2_cls['macro_averaged'] = np.array([v for k, v in psds2_cls.items()]).mean()
+    # psds2_cls['name'] = 'psds2'
+    psds_cls = pd.DataFrame([psds1_cls, psds1_cls_tsed])
+    # psds_cls = pd.DataFrame([psds1_cls, psds2_cls])
+    os.makedirs(f'{save_path}/psds_cls/', exist_ok=True)
+    psds_cls.to_csv(f'{save_path}/psds_cls/{epoch}.csv', index=False)
+
+    return psds1_sed_scores_eval, psds1_sed_scores_eval_tsed

src/clap_embedding/Accelerating, revving, vroom.pt

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

src/clap_embedding/Air brake.pt

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

src/clap_embedding/Air conditioning.pt

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

src/clap_embedding/Air horn, truck horn.pt

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

src/clap_embedding/Aircraft engine.pt

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

src/clap_embedding/Aircraft.pt

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

src/clap_embedding/Alarm clock.pt

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

src/clap_embedding/Alarm.pt

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

src/clap_embedding/Alert.pt

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

src/clap_embedding/Ambulance (siren).pt

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

src/clap_embedding/Animal.pt

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

src/clap_embedding/Applause.pt

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

src/clap_embedding/Arrow.pt

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

src/clap_embedding/Artillery fire.pt

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

src/clap_embedding/Audio logo.pt

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

src/clap_embedding/Babbling.pt

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

src/clap_embedding/Baby cry, infant cry.pt

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

src/clap_embedding/Baby laughter.pt

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

src/clap_embedding/Background noise.pt

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

src/clap_embedding/Bang.pt

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

src/clap_embedding/Bark.pt

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

src/clap_embedding/Basketball bounce.pt

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

src/clap_embedding/Bathroom sounds.pt

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

src/clap_embedding/Bathtub (filling or washing).pt

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

src/clap_embedding/Battle cry.pt

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

src/clap_embedding/Bee, wasp, etc..pt

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

src/clap_embedding/Beep, bleep.pt

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

src/clap_embedding/Bell.pt

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

src/clap_embedding/Bellow.pt

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

src/clap_embedding/Belly laugh.pt

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

src/clap_embedding/Bicycle bell.pt

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

src/clap_embedding/Bicycle, tricycle.pt

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

src/clap_embedding/Bird flight, flapping wings.pt

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

src/clap_embedding/Bird vocalization, bird call, bird song.pt

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

src/clap_embedding/Bird.pt

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

src/clap_embedding/Biting.pt

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

src/clap_embedding/Bleat.pt

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

src/clap_embedding/Blender, food processor.pt

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