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