Skip to content

Edge AI solution for real-time Wet/Dry cough classification on ESP32 using a specialized TinyML model (87.6% overall accuracy).

Notifications You must be signed in to change notification settings

mnouira02/CoughGuard-TinyML

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

1 Commit
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

🛡️ CoughGuard: TinyML Cough Classification on ESP32

CoughGuard is an end-to-end TinyML application that detects and classifies coughs ("Wet" vs "Dry") in real-time on an ESP32 microcontroller using a custom Convolutional Neural Network (CNN).

Wokwi Simulation

🚀 Overview

Respiratory conditions often require monitoring the type of cough a patient has. This project deploys a lightweight AI model directly onto an ESP32 to classify audio in real-time, privacy-preserving, and offline.

  • Input: Raw audio (processed into Mel-Spectrograms).
  • Output: Classification ("Wet" vs. "Dry") on an OLED Display.
  • Platform: ESP32 (Arduino Framework) + TensorFlow Lite Micro.

🎧 Audio Demonstration

The core challenge of this project is distinguishing the subtle differences between these two cough types.

Type Audio Sample
Dry Cough
Wet Cough

📊 Model Performance

The model was trained on a dataset of ~14,000 augmented samples and quantized to int8 for hardware deployment.

Metric Result Notes
Overall Accuracy 87.6% Tested on 9,341 held-out samples.
Dry Recall 94.8% Highly reliable at identifying normal coughs.
Wet Recall 64.7% Effective screening for pathological coughs.
Model Size ~60 KB Optimized for ESP32 RAM constraints.

Confusion Matrix

🧠 The Engineering Challenge

1. The Accuracy Paradox

The original dataset was heavily imbalanced (7,000 Dry vs. 2,000 Wet). Initial models achieved 80% accuracy by simply guessing "Dry" (ignoring sickness).

  • Solution: I implemented a custom training pipeline that augments the minority class (Noise injection + Pitch shifting) and downsamples the majority class to force a strict 50/50 training split.

2. Quantization Drift

Converting the model from Float32 (Python) to Int8 (ESP32) initially caused a "Negative 50%" confidence error due to mismatched Zero Points.

  • Solution: Implemented Dynamic Quantization logic in the evaluation script to read the specific scale and zero_point parameters from the TFLite interpreter, ensuring the Python evaluation matches the hardware inference exactly.

🛠️ Tech Stack

  • Hardware: ESP32 DevKit V4, SSD1306 OLED Display.
  • Firmware: C++, Arduino Framework, TensorFlow Lite for Microcontrollers.
  • ML Pipeline: TensorFlow/Keras, Librosa (Audio Processing), Pandas.
  • Tools: Streamlit (Data Auditing), Wokwi (Hardware Simulation).

📂 Project Structure

CoughGuard/
├── dataset/                   # Raw audio files (Dry/Wet)
├── firmware/                  # ESP32 C++ Code
│   ├── src/
│   │   ├── main.cpp           # Inference logic & OLED control
│   │   ├── model_data.h       # Auto-generated TFLite model (Hex)
│   │   └── test_samples.h     # Auto-generated test audio clips
│   └── platformio.ini
├── scripts/                   # Python ML Pipeline
│   ├── train_and_export.py    # Training, Augmentation & Quantization
│   ├── evaluate_model.py      # Accuracy testing & Confusion Matrix
│   ├── generate_test_header.py# Generates C++ headers for simulation
│   └── cough_review_app.py    # Streamlit tool for dataset cleaning
├── images/                    # Screenshots for README
└── requirements.txt

⚡ Getting Started

0. Acquire Dataset

To run the Python training and evaluation scripts, you must first download the expert-labeled subset of the COUGHVID dataset.

  1. Download the required raw audio files from COUGHVID on Zenodo.
  2. Unzip the files.
  3. Create the directory structure dataset/dry and dataset/wet and place the respective audio files inside.

Without these files, the training script will not be able to perform augmentation and balancing.

1. Setup Environment

Clone the repo and install dependencies:

pip install -r requirements.txt

2. Train the Model

This script balances the data, trains the CNN, quantizes it to int8, and exports the model_data.h C++ file directly to the firmware folder.

python scripts/train_and_export.py

3. Evaluate Performance

Generates the accuracy report and confusion matrix based on the quantized TFLite model.

python scripts/evaluate_model.py

4. Deploy / Simulate

  • Simulation: Use Wokwi (link to your project).
  • Physical: Open the firmware/ folder in PlatformIO or Arduino IDE and flash to your ESP32.

📚 Dataset Reference

This project uses the COUGHVID crowdsourcing dataset, specifically the subset annotated by expert pulmonologists.

  • Dataset Source: COUGHVID on Zenodo
  • Original Paper: The COUGHVID crowdsourcing dataset, a corpus for the study of large-scale cough analysis algorithms (Orlandic et al., 2021).

Citation:

Orlandic, L., Teijeiro, T. & Atienza, D. The COUGHVID crowdsourcing dataset, a corpus for the study of large-scale cough analysis algorithms. Sci Data 8, 156 (2021). https://doi.org/10.1038/s41597-021-00937-4

👨‍💻 Author

Marouane Nouira

📜 License

MIT

About

Edge AI solution for real-time Wet/Dry cough classification on ESP32 using a specialized TinyML model (87.6% overall accuracy).

Topics

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published