joakes90's picture
Clean up repo: proper model card, drop Space-specific gradio_app/ contents
e6ff95c verified
|
Raw
History Blame Contribute Delete
4.17 kB
metadata
license: cc-by-sa-4.0
pipeline_tag: audio-classification
tags:
  - audio-classification
  - audio
  - pytorch
  - panns
  - audioset
  - engine-sound

Engine Sound Classifier (PANNs CNN14 fine-tune)

A fine-tuned PANNs CNN14 (pretrained on AudioSet) that classifies short audio clips of running engines into one of 34 configurations — cylinder count, layout, and stroke type (e.g. i4, v8_cross, single_two_stroke, 2_rotor).

Status: early, ambitious, not yet accurate. This is a first published checkpoint, shared to show the approach and current state, warts included. Recording-level balanced accuracy is ~0.33 on held-out test data (random guessing over 34 classes is ~0.03).

An informational demo Space (sample clips, performance breakdown) is at joakes90/engine-sound-classifier. The training code, notebook, and full diagnostic write-up live at github.com/joakes90/auto_sound_train.

Files

  • best_model_run_d_soft_weights.pth — fine-tuned state dict (~325 MB), Run D checkpoint (best balanced accuracy so far).
  • class_names.json — ordered list of the 34 class labels the output layer corresponds to.

Model details

  • Architecture: CNN14 from PANNs, pretrained on AudioSet, with the final fc_audioset layer replaced by a 34-way linear head and fine-tuned.
  • Input: mono audio, resampled to 32 kHz, in 2-second windows (matching the training manifest's window length).
  • Output: softmax over 34 engine configurations. For clips longer than 2s, run overlapping windows (e.g. 1s hop) and average per-window probabilities before taking the top class — this "recording-level" pooling is consistently more accurate than scoring a single window.

Classes

2_rotor, h12, h2, h4, h6, i2_180, i2_180_two_stroke, i2_270, i2_360, i2_360_two_stroke, i3, i4, i4_crossplane, i4_diesel, i5, i5_diesel, i6, i6_diesel, single_four_stroke, single_two_stroke, v10_72, v12, v2_45, v2_90, v4, v4_two_stroke, v6_120, v6_60, v6_90_even, v8_cross, v8_diesel, v8_flat, v8_voodoo, vr6

Performance

Run Change Balanced acc Micro acc
A baseline (no regularization) — 0.38 (overfit)
B + SpecAugment/noise/mixup 0.297 0.311 (underfit)
C fixed SpecAugment time-mask scale 0.323 0.343
D + softened class weights (counts^-0.5) 0.330 0.376

This checkpoint is Run D. Cylinder-family accuracy (does it at least get the cylinder count right) is meaningfully higher than exact-class accuracy — most confusion is between siblings within the same engine family, not wild misfires. Known weak spots: engines with many cylinders (8/12) are confused with close siblings far more than 1–2 cylinder engines, and a couple of classes (v12, 2_rotor) remain poorly calibrated. Full diagnosis in FINDINGS.md in the training repo.

Usage

import json
import torch
from huggingface_hub import hf_hub_download
from panns_inference import AudioTagging

REPO_ID = "joakes90/engine_sound_cassifier"
CHECKPOINT = "best_model_run_d_soft_weights.pth"

class_names = json.load(
    open(hf_hub_download(REPO_ID, "class_names.json"))
)["class_names"]

# Loads PANNs' pretrained AudioSet CNN14, then swaps in our fine-tuned head.
at = AudioTagging(checkpoint_path=None, device="cpu")
backbone = at.model.module if isinstance(at.model, torch.nn.DataParallel) else at.model
backbone.fc_audioset = torch.nn.Linear(backbone.fc_audioset.in_features, len(class_names))

state_dict = torch.load(hf_hub_download(REPO_ID, CHECKPOINT), map_location="cpu")
backbone.load_state_dict(state_dict)
backbone.eval()

# waveform: mono float tensor at 32kHz, shape (num_windows, window_samples)
with torch.inference_mode():
    logits = backbone.fc_audioset(backbone(waveform)["embedding"])
    probs = torch.softmax(logits, dim=1).mean(dim=0)
    top_class = class_names[probs.argmax()]

License

CC BY-SA 4.0.