| import numpy as np
|
| import torchaudio
|
| import torchaudio.transforms as T
|
| import joblib
|
| from scipy.stats import skew, kurtosis
|
| import tensorflow_hub as hub
|
|
|
|
|
| clf = joblib.load("models/noise_classifier.pkl")
|
| label_encoder = joblib.load("models/label_encoder.pkl")
|
|
|
|
|
| yamnet_model = hub.load("https://tfhub.dev/google/yamnet/1")
|
|
|
| def get_yamnet_embedding(audio_path):
|
| """
|
| Extract YAMNet embeddings with statistical pooling from a WAV file.
|
| """
|
| try:
|
| waveform, sr = torchaudio.load(audio_path)
|
| if sr != 16000:
|
| resampler = T.Resample(orig_freq=sr, new_freq=16000)
|
| waveform = resampler(waveform)
|
| if waveform.size(0) > 1:
|
| waveform = waveform.mean(dim=0)
|
| else:
|
| waveform = waveform.squeeze(0)
|
|
|
| waveform_np = waveform.numpy()
|
| _, embeddings, _ = yamnet_model(waveform_np)
|
|
|
|
|
| mean = np.mean(embeddings, axis=0)
|
| std = np.std(embeddings, axis=0)
|
| min_val = np.min(embeddings, axis=0)
|
| max_val = np.max(embeddings, axis=0)
|
| skewness = skew(embeddings, axis=0)
|
| kurt = kurtosis(embeddings, axis=0)
|
|
|
| return np.concatenate([mean, std, min_val, max_val, skewness, kurt])
|
| except Exception as e:
|
| print(f"Failed to process {audio_path}: {e}")
|
| return None
|
|
|
| def classify_noise(audio_path, threshold=0.6):
|
| """
|
| Classify noise with rejection threshold for 'Unknown' label.
|
| """
|
| feature = get_yamnet_embedding(audio_path)
|
| if feature is None:
|
| return [("Unknown", 0.0)]
|
|
|
| feature = feature.reshape(1, -1)
|
| probs = clf.predict_proba(feature)[0]
|
|
|
| top_idx = np.argmax(probs)
|
| top_prob = probs[top_idx]
|
|
|
| if top_prob < threshold:
|
| return [("Unknown", top_prob)]
|
|
|
| top_indices = np.argsort(probs)[::-1][:5]
|
| return [(label_encoder.inverse_transform([i])[0], probs[i]) for i in top_indices]
|
|
|
