Spaces:
Runtime error
Runtime error
File size: 1,260 Bytes
b7becdf | 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 | #!/usr/bin/env python3
import os
from sound_classifier import SoundClassifier
def train_and_save_models(data_dir='data', models_dir='models'):
"""
Train and save multiple models for engine sound classification.
Args:
data_dir (str): Directory containing the sound data
models_dir (str): Directory to save the trained models
"""
# Ensure models directory exists
os.makedirs(models_dir, exist_ok=True)
# Model types to train
model_types = ['rf', 'lr', 'svm', 'nn']
for model_type in model_types:
print(f"\n{'='*50}")
print(f"Training {model_type.upper()} model with benchmark data as 'normal'...")
print(f"{'='*50}")
# Initialize classifier with benchmark data included
classifier = SoundClassifier(
data_dir=data_dir,
model_type=model_type,
include_benchmark=True
)
# Train the model
classifier.train()
# Save the model
model_path = os.path.join(models_dir, f"{model_type}_sound_classifier_model.joblib")
classifier.save_model(model_path)
print(f"Model saved to {model_path}")
if __name__ == "__main__":
train_and_save_models()
|