File size: 4,167 Bytes
693dda4
 
e6ff95c
 
 
 
 
 
 
 
693dda4
e6ff95c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
---
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](https://github.com/qiuqiangkong/audioset_tagging_cnn)
(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](https://huggingface.co/spaces/joakes90/engine-sound-classifier).
The training code, notebook, and full diagnostic write-up live at
[github.com/joakes90/auto_sound_train](https://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

```python
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.