wav2vec2-base-drum-kit

Wav2Vec2 for drum-kit classification

Fine-tuned facebook/wav2vec2-base for audio classification of single drum/percussion sounds into 10 classes.

Classes

• clap, conga, crash, cymbal, hat, kick, ride, rim, snare, tom

Usage

import torch
import librosa
from transformers import AutoFeatureExtractor, AutoModelForAudioClassification

model_id = "airasoul/wav2vec2-base-drum-kit"  # e.g. username/wav2vec2-base-drum-kit
feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
model = AutoModelForAudioClassification.from_pretrained(model_id)
model.eval()

# Load a WAV (16 kHz mono)
audio, sr = librosa.load("path/to/audio.wav", sr=16000, mono=True)
inputs = feature_extractor(
    audio, sampling_rate=16000, max_length=48000,  # 3 s at 16 kHz
    truncation=True, return_tensors="pt", padding=True
)
with torch.no_grad():
    logits = model(**inputs).logits
pred_id = logits.argmax(dim=-1).item()
label = model.config.id2label.get(pred_id) or model.config.id2label.get(str(pred_id))
print(label)  # e.g. "kick"

Training

• Base: facebook/wav2vec2-base
• Task: Single-label classification over 10 drum classes
• Data: Custom drum-kit dataset with augmentation (time stretch, noise, gain)
• Input: 16 kHz mono, up to 3 s (truncated or padded)

Results

• Validation accuracy: 95.7% (epoch 10)
• Test accuracy: 97.0% (300 samples, held-out)

Limitations

• Trained on short, single-hit drum sounds. Performance may drop on long mixes, multiple overlapping sounds, or very different recording conditions.