main.py

import os

import torch
import beats.BEATs
from beats.BEATs import BEATs, BEATsConfig
import csv
import soundfile as sf
import librosa
import numpy as np

BEATs_model: beats.BEATs.BEATs
label_mapping = {}
checkpoint = {}

def load_beats_model():
    global BEATs_model, checkpoint, label_mapping
    # load the pre-trained checkpoints
    checkpoint = torch.load('beats/BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt2.pt')

    # load the labels
    with open("data/class_labels_indices.csv", mode="r") as file:
        reader = csv.reader(file)
        for row in reader:
            label_mapping[row[1]] = row[2]

    # load the model
    cfg = BEATsConfig(checkpoint['cfg'])
    BEATs_model = BEATs(cfg)
    BEATs_model.load_state_dict(checkpoint['model'])
    BEATs_model.eval()

def beats_predict_labels(audio):
    with torch.inference_mode():
        probs = BEATs_model.extract_features(torch.tensor(audio, dtype=torch.float32))[0]

    label_prob, label_idx = probs.sort()
    label_name = [[label_mapping[checkpoint['label_dict'][idx.item()]] for idx in row] for row in label_idx][0]
    label_idx = [[checkpoint['label_dict'][idx.item()] for idx in row] for row in label_idx][0]

    return label_idx, label_name, label_prob[0]

def stream_audio_chunks(path, chunk_length_seconds, stride_seconds=1.0, target_sr=16000):
    with sf.SoundFile(path) as f:
        orig_sr = f.samplerate

        chunk_size = int(chunk_length_seconds * orig_sr)
        stride_size = int(stride_seconds * orig_sr)

        buffer = f.read(chunk_size)

        if buffer.ndim == 2:
            buffer = buffer.mean(axis=1)

        while len(buffer) == chunk_size:
            # Resample
            chunk = librosa.resample(buffer, orig_sr=orig_sr, target_sr=target_sr)
            yield chunk

            # Move forward by stride
            f.seek(f.tell() - (chunk_size - stride_size))

            buffer = f.read(chunk_size)

            if buffer.ndim == 2:
                buffer = buffer.mean(axis=1)

def get_label_probs(audio_path):
    label_probs = {}

    for chunk in stream_audio_chunks(audio_path, 5):
        label_idxs, label_names, probs = beats_predict_labels(chunk[np.newaxis, :])
        for idx, name, prob in zip(label_idxs, label_names, probs):
            if idx not in label_probs:
                label_probs[idx] = {"probs": []}
            label_probs[idx]["display_name"] = name
            label_probs[idx]["probs"].append(prob.item())

    return label_probs

# Get predicted labels from the specified audio
def get_audio_summary(audio_path):
    summary = {}
    label_probs = get_label_probs(audio_path=audio_path)
    for idx, d in label_probs.items():
        r = 2
        avg_conf = pow(sum(np.pow(d['probs'],r))/len(d['probs']), 1/r)

        summary[idx] = {
            "display_name": d["display_name"],
            "avg_confidence": avg_conf
        }
    return summary


def generate_images(path):
    files = []
    if os.path.isdir(path):
        for file in os.listdir(path):
            audio_file = os.path.join(path, file)
            files.append(audio_file)
    else:
        files.append(path)

    for audio_file in files:
        print(f"Generating audio summary for {audio_file}...")
        summary = get_audio_summary(audio_file)

        print(sorted(summary.items(), key=lambda x: x[1]["avg_confidence"], reverse=True)[:10])

        template = UrbanSquareTemplate("data/audioset_parent_mapping.csv")

        labels, probs = zip(*[(str(k).lower(), v["avg_confidence"]) for k, v in summary.items()])

        file_name = os.path.splitext(audio_file.split("/")[-1])[0]
        template.generate_env(labels, probs, out_path=f"./env_templates/urban_square/out2/{file_name}.png")


if __name__ == "__main__":
    print("Loading BEATs model...")
    load_beats_model()

    from env_templates.urban_square.urban_square import UrbanSquareTemplate

    _dir = "./test_sounds/city_sounds/sound_excerpts/"
    generate_images(_dir)