Audio-Mambaを使った音声分類

import os
import torch
import torchaudio
import numpy as np
import src.models as models
from src import dataloader
from src.utilities.stats import calculate_stats
from IPython.display import Audio, display
import csv
import warnings

class Namespace:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)

data_args = Namespace(
    num_mel_bins = 128,
    target_length = 1024,
    mean = -5.0767093,
    std = 4.4533687,
)

model_args = Namespace(
    model_type = 'base',
    n_classes = 309,
    imagenet_pretrain = False,
    imagenet_pretrain_path = None,
    aum_pretrain = True,
    aum_pretrain_path = '{GoogleDriveからダウンロードしてきたモデルのPATH}',
    aum_variant = 'Fo-Bi',
    device = 'cuda',
)

AuM = models.AudioMamba(
    spectrogram_size=(data_args.num_mel_bins, data_args.target_length),
    patch_size=(16, 16),
    strides=(16, 16),
    embed_dim=768,
    num_classes=model_args.n_classes,
    imagenet_pretrain=model_args.imagenet_pretrain,
    imagenet_pretrain_path=model_args.imagenet_pretrain_path,
    aum_pretrain=model_args.aum_pretrain,
    aum_pretrain_path=model_args.aum_pretrain_path,
    bimamba_type='v1',
)

AuM.to(model_args.device)
AuM.eval()

label_dict = {}
with open("./Audio-Mamba-AuM/exps/vggsound/data/class_labels_indices.csv", 'r') as f:
    csv_reader = csv.DictReader(f)
    line_count = 0
    for row in csv_reader:
        label_dict[row['index']] = row['display_name']
        line_count += 1

def predict(audio_path):
    waveform, sr = torchaudio.load(audio_path)
    waveform = waveform - waveform.mean()

    fbank = torchaudio.compliance.kaldi.fbank(
        waveform, 
        htk_compat=True, 
        sample_frequency=sr, 
        use_energy=False,
        window_type='hanning', 
        num_mel_bins=data_args.num_mel_bins,
        dither=0.0, 
        frame_shift=10
    )

    n_frames = fbank.shape[0]
    p = data_args.target_length - n_frames

    if p > 0:
        m = torch.nn.ZeroPad2d((0, 0, 0, p))
        fbank = m(fbank)
    elif p < 0:
        fbank = fbank[0:data_args.target_length, :]

    label_indices = np.zeros(model_args.n_classes)
    label_indices = torch.FloatTensor(label_indices)

    fbank = (fbank - data_args.mean) / (data_args.std * 2)

    fbank = fbank.unsqueeze(0)
    label_indices = label_indices.unsqueeze(0)

    fbank = fbank.to(model_args.device)
    label_indices = label_indices.to(model_args.device)

    with torch.no_grad():
        output = AuM(fbank)

    output = torch.sigmoid(output)
    output = output.cpu().numpy()

    top_idx = np.argsort(output[0])[-1]
    return label_dict[str(top_idx)]