Spaces:
Runtime error
Runtime error
File size: 5,050 Bytes
2c2f868 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
import os
import numpy as np
import librosa
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import classification_report
import joblib
class SoundClassifier:
def __init__(self, data_dir, model_type='rf', sr=22050, duration=20):
self.data_dir = data_dir
self.sr = sr
self.duration = duration
self.model = None
self.le = LabelEncoder()
self.scaler = StandardScaler()
self.model_type = model_type
def extract_features(self, file_path):
# Load audio file
y, _ = librosa.load(file_path, sr=self.sr, duration=self.duration)
# Pad or truncate to fixed length
if len(y) < self.sr * self.duration:
y = np.pad(y, (0, self.sr * self.duration - len(y)))
else:
y = y[:self.sr * self.duration]
# Extract features
mfccs = librosa.feature.mfcc(y=y, sr=self.sr, n_mfcc=13)
spectral_centroid = librosa.feature.spectral_centroid(y=y, sr=self.sr)
spectral_rolloff = librosa.feature.spectral_rolloff(y=y, sr=self.sr)
# Compute statistics
features = np.concatenate([
mfccs.mean(axis=1),
mfccs.std(axis=1),
spectral_centroid.mean(axis=1),
spectral_rolloff.mean(axis=1)
])
return features
def prepare_data(self):
X = []
y = []
# Iterate through each issue folder
for issue in os.listdir(self.data_dir):
issue_path = os.path.join(self.data_dir, issue)
if os.path.isdir(issue_path):
# Process each audio file in the folder
for audio_file in os.listdir(issue_path):
if audio_file.endswith('.wav'):
file_path = os.path.join(issue_path, audio_file)
features = self.extract_features(file_path)
X.append(features)
y.append(issue)
print(len(X))
print(len(y))
X = np.array(X)
y = self.le.fit_transform(y)
return X, y
def train(self):
# Prepare data
X, y = self.prepare_data()
# Split data
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42
)
# Scale features
X_train_scaled = self.scaler.fit_transform(X_train)
X_test_scaled = self.scaler.transform(X_test)
# Train model based on model_type
if self.model_type == 'rf':
self.model = RandomForestClassifier(n_estimators=100, random_state=42)
elif self.model_type == 'lr':
self.model = LogisticRegression(random_state=42, max_iter=1000)
elif self.model_type == 'svm':
self.model = SVC(kernel='rbf', random_state=42)
elif self.model_type == 'nn':
self.model = MLPClassifier(hidden_layer_sizes=(100, 50), max_iter=1000, random_state=42)
else:
raise ValueError("Invalid model type. Choose 'rf', 'lr', 'svm', or 'nn'.")
self.model.fit(X_train_scaled, y_train)
# Evaluate
y_pred = self.model.predict(X_test_scaled)
print(f"\nModel Performance ({self.model_type}):")
print(classification_report(y_test, y_pred,
labels=np.unique(y),
target_names=self.le.classes_[np.unique(y)]))
return self.model
def predict(self, audio_file):
# Extract features from new audio
features = self.extract_features(audio_file)
# Scale features
features_scaled = self.scaler.transform([features])
# Make prediction
prediction = self.model.predict(features_scaled)[0]
# Return the issue name
return self.le.inverse_transform([prediction])[0]
def save_model(self, model_path='sound_classifier_model.joblib'):
"""Save the trained model, label encoder, and scaler"""
if self.model is None:
raise ValueError("Model hasn't been trained yet!")
model_data = {
'model': self.model,
'label_encoder': self.le,
'scaler': self.scaler,
'model_type': self.model_type
}
joblib.dump(model_data, model_path)
@classmethod
def load_model(cls, model_path='sound_classifier_model.joblib'):
"""Load a trained model"""
classifier = cls(data_dir=None) # Create instance without data dir
model_data = joblib.load(model_path)
classifier.model = model_data['model']
classifier.le = model_data['label_encoder']
classifier.scaler = model_data['scaler']
classifier.model_type = model_data['model_type']
return classifier
|