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

Download model.py, predict.py and weightsfilefromspectrogram.pth into one folder, install the dependencies, then run:

pip install torch pillow numpy
Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support