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