File size: 2,723 Bytes
3d90862 bdb1abb ed2020b bdb1abb ed2020b b649a34 ed2020b 6a98a26 3d90862 ed2020b 04b5c34 ed2020b 3d90862 04b5c34 ed2020b 3d90862 ed2020b 04b5c34 ed2020b | 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 | ---
license: mit
library_name: pytorch
tags:
- audio-classification
- spectrogram
- cnn
- pytorch
---
# Drone / Helicopter / Background Audio Classifier
A small convolutional neural network, written and trained from scratch in PyTorch, that classifies spectrogram images of audio into three classes: **background**, **helicopter** and **drone**. Developed from work at Ministry of Defence hackathons (Sandhurst) and deployed independently in 2026.
## Classes (output index order)
| Index | Class |
|---|---|
| 0 | background |
| 1 | helicopter |
| 2 | drone |
## Model
- **Input:** single-channel (grayscale) spectrogram image, 480 x 640 pixels, pixel values scaled to 0-1.
- **Architecture:** two convolutional layers (64 then 128 filters, 3x3 kernels), each followed by ReLU and 2x2 max pooling, then one linear layer to 3 classes. Output is log-probabilities (`log_softmax`), so apply `exp` to get probabilities.
- **Code:** the architecture is in `model.py` (`CnnNet`) and a minimal inference example is in `predict.py`.
- **Weights:** `weightsfilefromspectrogram.pth` (PyTorch `state_dict`).
## Data
- 297 spectrogram images: 101 background, 99 helicopter, 97 drone.
- Spectrograms were generated from dataset provided by the London Defence Hackathon community on GitHub
- Random 60/20/20 train/validation/test split (`random_state=42`), about 178 / 59 / 60 images.
## Training
- Framework: PyTorch, trained on Google Colab.
- Loss: negative log-likelihood. Optimiser: SGD, learning rate 0.001, momentum 0.5, batch size 4, 115 epochs.
- The checkpoint with the best validation accuracy was kept.
## Results (60 held-out test spectrograms)
- Test accuracy: 86.67%
- Weighted F1: 0.87 (precision 0.87, recall 0.87)
- Per-class one-vs-rest ROC-AUC: helicopter 0.99, drone 0.96, background 0.94
## Limitations
- Small dataset of 297 spectrograms and a small test set (60 images) so these figures are indicative rather than precise.
- Training loss is very low while validation accuracy is lower, which suggests overfitting. No data augmentation or regularisation was used.
- Background clips are sometimes misclassified as helicopter.
- The split is random at image level. If several spectrograms come from the same recording, near-duplicates could appear in both train and test and inflate the scores; splitting by recording is a planned improvement.
- Not tested on recordings from new environments, microphones or noise conditions.
- This is a baseline model for problem-solving demonstration. Not for safety-critical use.
## Usage
```python
Download model.py, predict.py and weightsfilefromspectrogram.pth into one folder, install the dependencies, then run:
pip install torch pillow numpy
``` |