Zen-4011/audio-separation-model

1import os2import torch3import torchaudio4import pandas as pd5import numpy as np6import librosa7import librosa.display8import matplotlib9matplotlib.use('Agg')  # prevents matplotlib from trying to open a GUI window10import matplotlib.pyplot as plt11from Core.resnet_model import AudioResNet12from Core.gtzan_dataset import GENRES13 14device = torch.device("cpu")15 16# Confidence Threshold (< .20%)17NATURE_CONFIDENCE_THRESHOLD = 0.2018 19# mel Spectrogram transform20mel_transform = torchaudio.transforms.MelSpectrogram(21    sample_rate = 22050,22    n_fft = 1024,23    hop_length = 512,24    n_mels = 12825).to(device)26 27# ESC-50 class map 28def load_esc50_classes(csv_path = "data/esc50.csv"):29    df = pd.read_csv(csv_path)30    class_map = dict(zip(df['target'], df['category']))31    return class_map32 33# Model loader 34def load_model(num_classes, weights_path):35    if not os.path.exists(weights_path):36        print(f"Warning: weights not found at {weights_path}")37        return None38    model = AudioResNet(num_classes = num_classes).to(device)39    model.load_state_dict(torch.load(40        weights_path, map_location = device, weights_only = True41    ))42    model.eval()43    print(f"Loaded: {weights_path}")44    return model45 46# Load both models at startup47nature_model = load_model(num_classes = 50, weights_path = "Models/esc50_resnet_v1.pth")48music_model  = load_model(num_classes = 10, weights_path = "Models/gtzan_resnet_v1.pth")49 50try:51    ESC50_CLASSES = load_esc50_classes()52except FileNotFoundError:53    print("Warning: esc50.csv not found.")54    ESC50_CLASSES = {}55 56def models_are_loaded():57    return nature_model is not None and music_model is not None58 59# Audio preprocessor 60def preprocess_audio(audio_path, num_samples):61    signal, sr = torchaudio.load(audio_path)62 63    if sr != 22050:64        signal = torchaudio.transforms.Resample(sr, 22050)(signal)65 66    if signal.shape[0] > 1:67        signal = torch.mean(signal, dim = 0, keepdim = True)68 69    if signal.shape[1] > num_samples:70        signal = signal[:, :num_samples]71    elif signal.shape[1] < num_samples:72        signal = torch.nn.functional.pad(signal, (0, num_samples - signal.shape[1]))73 74    signal = signal.to(device)75    mel    = mel_transform(signal).unsqueeze(0)76    return mel77 78# Nature prediction —> returns top 3 + recognised flag79def predict_nature(audio_path):80    """81    Returns a dict with:82      - recognised (bool)83      - label (str)         — top prediction, or "Unrecognised Sound"84      - closest_match (str) — always the top prediction regardless of threshold85      - confidence (float)  — top prediction confidence %86      - top3 (list)         — [{label, confidence}, ...] always 3 items87    """88    if nature_model is None:89        return {90            "recognised":    False,91            "label":         "Model not loaded",92            "closest_match": "Model not loaded",93            "confidence":    0.0,94            "top3":          []95        }96 97    mel = preprocess_audio(audio_path, num_samples=22050 * 5)98 99    with torch.no_grad():100        outputs       = nature_model(mel)101        probabilities = torch.nn.functional.softmax(outputs / 3.0, dim=1)102 103        # Top 3 predictions104        top3_confidences, top3_indices = torch.topk(probabilities, k=3, dim=1)105 106    top3 = []107    for i in range(3):108        idx        = top3_indices[0][i].item()109        conf       = round(top3_confidences[0][i].item(), 4)110        raw_label  = ESC50_CLASSES.get(idx, "Unknown")111        clean_label = raw_label.replace('_', ' ').title()112        top3.append({"label": clean_label, "confidence": conf})113 114    top_label      = top3[0]["label"]115    top_confidence = top3[0]["confidence"]116    recognised     = top_confidence >= 0.25117 118    return {119        "recognised":    recognised,120        "label":         top_label if recognised else "Unrecognised Sound",121        "closest_match": top_label,122        "confidence":    top_confidence,123        "top3":          top3124    }125 126# Music prediction —> returns top 3 + recognised flag127def predict_music(audio_path):128    """129    Returns a dict with:130      - recognised (bool)131      - label (str)132      - closest_match (str)133      - confidence (float)134      - top3 (list)135    """136    if music_model is None:137        return {138            "recognised":    False,139            "label":         "Model not loaded",140            "closest_match": "Model not loaded",141            "confidence":    0.0,142            "top3":          []143        }144 145    mel = preprocess_audio(audio_path, num_samples = 22050 * 30)146 147    with torch.no_grad():148        outputs       = music_model(mel)149        probabilities = torch.nn.functional.softmax(outputs, dim = 1)150 151        top3_confidences, top3_indices = torch.topk(probabilities, k = 3, dim = 1)152 153    top3 = []154    for i in range(3):155        idx   = top3_indices[0][i].item()156        conf  = round(top3_confidences[0][i].item(), 4)157        label = GENRES[idx].title()158        top3.append({"label": label, "confidence": conf})159 160    top_label      = top3[0]["label"]161    top_confidence = top3[0]["confidence"]162    recognised     = top_confidence >= 0.25163 164    return {165        "recognised":    recognised,166        "label":         top_label if recognised else "Unrecognised Sound",167        "closest_match": top_label,168        "confidence":    top_confidence,169        "top3":          top3170    }171 172# Spectrogram image generator173def generate_spectrogram_image(audio_path, save_path, title=None):174    """175    Generates a styled mel spectrogram image for a given audio stem.176    Saves to save_path and returns the path.177    Uses the magma colormap — looks great on dark-themed frontends.178    """179    try:180        y, sr = librosa.load(audio_path, sr = 22050)181 182        mel    = librosa.feature.melspectrogram(183            y = y, 184            sr = sr, 185            n_fft = 1024, 186            hop_length = 512, 187            n_mels = 128188        )189 190        mel_db = librosa.power_to_db(mel, ref = np.max)191 192        fig, ax = plt.subplots(figsize = (8, 3), facecolor = '#1a1a2e')193        ax.set_facecolor('#1a1a2e')194 195        img = librosa.display.specshow(196            mel_db,197            sr = sr,198            hop_length = 512,199            x_axis = 'time',200            y_axis = 'mel',201            cmap = 'magma',202            ax = ax203        )204 205        cbar = fig.colorbar(img, ax = ax, format = '%+2.0f dB')206        cbar.ax.yaxis.set_tick_params(color = 'white')207        plt.setp(cbar.ax.yaxis.get_ticklabels(), color = 'white', fontsize = 8)208 209        display_title = title or os.path.basename(audio_path).replace('.wav', '').title()210        ax.set_title(display_title, color = 'white', fontsize = 12, fontweight = 'bold', pad = 8)211        ax.tick_params(colors = 'white', labelsize = 8)212        ax.xaxis.label.set_color('white')213        ax.yaxis.label.set_color('white')214 215        for spine in ax.spines.values():216            spine.set_edgecolor('#444444')217 218        plt.tight_layout()219 220        os.makedirs(os.path.dirname(save_path) if os.path.dirname(save_path) else '.', exist_ok = True)221        plt.savefig(save_path, dpi = 120, bbox_inches = 'tight', facecolor = '#1a1a2e')222        plt.close(fig)223 224        return save_path225 226    except Exception as e:227        print(f"Spectrogram generation failed for {audio_path}: {e}")228        plt.close('all')229        return None