Armak/SED

1from pathlib import Path2from typing import Tuple3 4import dcase_util5import matplotlib.pyplot as plt6import numpy as np7import pandas as pd8import sed_vis9import torch10import yaml11 12from desed_task.dataio.datasets import StronglyAnnotatedSet13from encoder import ManyHotEncoder14from model import CRNN15from trainer import SED16from utils import batched_decode_preds, classes_labels17 18 19def load_config():20    """21    Loads the configuration parameters from a yaml file.22 23    Returns:24        dict: The configuration parameters.25    """26    with open("params.yaml", "r") as f:27        config = yaml.safe_load(f)28    return config29 30 31def initialize_sed(config: dict) -> Tuple[SED, ManyHotEncoder]:32    """33    Initializes the Sound Event Detection (SED) model and the encoder.34 35    Parameters:36        config (dict): The configuration parameters.37 38    Returns:39        tuple: The initialized SED model and encoder.40    """41    encoder = ManyHotEncoder(42        list(classes_labels.keys()),43        audio_len=config["data"]["audio_max_len"],44        frame_len=config["feats"]["n_filters"],45        frame_hop=config["feats"]["hop_length"],46        net_pooling=config["data"]["net_subsample"],47        fs=config["data"]["fs"],48    )49 50    sed = SED(config, encoder=encoder, sed=CRNN(**config["net"]))51    return sed, encoder52 53 54def load_model(55    sed: SED, model_path: str = "../dvclive/artifacts/epoch=68-step=7314.ckpt"56) -> SED:57    """58    Loads a pretrained SED model from a checkpoint file.59 60    Parameters:61        sed (SED): The SED model to load the weights into.62        model_path (str): The path to the checkpoint file.63 64    Returns:65        SED: The SED model with loaded weights.66    """67    state_dict = torch.load(model_path, map_location=torch.device("cpu"))["state_dict"]68    sed.load_state_dict(state_dict)69    return sed70 71 72def make_prediction_on_audio(73    sed: SED,74    audio: torch.Tensor,75    config: dict,76    encoder: ManyHotEncoder,77    filename: str,78    model: str = "student",79    median_filter: int = 7,80) -> Tuple[81    sed_vis.visualization.EventListVisualizer, dcase_util.containers.AudioContainer82]:83    """84    Makes a prediction on a provided audio tensor.85 86    Parameters:87        sed (SED): The SED model to use for prediction.88        audio (torch.Tensor): The audio data as a tensor.89        config (dict): The configuration parameters.90        encoder (ManyHotEncoder): The encoder for the dataset.91        filename (str): The filename associated with the audio tensor.92        print_df (bool, optional): Whether to print the dataframes. Defaults to False.93 94    Returns:95        tuple: The visualizer for the event list and the audio container.96    """97    preds = batched_decode_preds(98        sed.forward(audio, model=model), filename, encoder, median_filter=median_filter99    )100    preds[2][0.5][["onset", "offset", "event_label"]].to_csv(101        "preds.tsv", index=False, sep="\t"102    )103 104    # Create an AudioContainer from the audio tensor105    audio_container = dcase_util.containers.AudioContainer(106        data=audio.numpy(),107        fs=config["data"]["fs"],108    )109 110    # Load event lists111    estimated_event_list = dcase_util.containers.MetaDataContainer().load("preds.tsv")112 113    event_lists = {114        "estimated": estimated_event_list,115    }116 117    vis = sed_vis.visualization.EventListVisualizer(118        event_lists=event_lists,119        audio_signal=audio_container.data,120        sampling_rate=audio_container.fs,121        publication_mode=True,122    )123    return vis, audio_container124