Dasheng-RNN Drum Transcriber

An interactive web application for automatic drum transcription using a Dasheng-based RNN model. This tool allows users to upload an audio file of a drum performance and receive a detailed transcription as a MIDI file.

The project includes not only the interactive demo but also a complete framework for training, evaluating, and analyzing custom drum transcription models.

✨ Features

• Interactive Web UI: Built with Gradio for an easy-to-use experience.
• Model Selection: Dynamically scan a directory for trained models and load them on the fly.
• Automatic Transcription: Upload an audio file (.wav, .mp3) and get a full drum transcription.
• Multiple Output Formats:
  • Visual Piano Roll: An interactive, scrollable piano roll to inspect the transcription.
  • Synthesized Audio: Listen to the transcribed MIDI rendered into audio for direct comparison with the original.
  • MIDI File: Download the .mid file for use in any Digital Audio Workstation (DAW).
• Optional Post-Quantization: Align transcribed notes to a precise rhythmic grid by specifying BPM and subdivision.
• Complete MLOps Pipeline: Includes scripts for training, evaluation, and in-depth error analysis.

🛠️ Tech Stack & Architecture

The application is built upon a modular architecture with several key components:

• UI Framework: Gradio
• Machine Learning Framework: PyTorch
• Audio Processing: Librosa
• MIDI Handling: PrettyMIDI
• Audio Synthesis: FluidSynth (with the GeneralUser GS soundfont)

The core transcription model is a Recurrent Neural Network (RNN) head built on top of the pre-trained Dasheng feature extractor.

📁 Project Structure

The codebase is organized to separate data handling, model logic, and application code.

.
├── dasheng_rnn/
│   ├── datasets/               # PyTorch Dataset classes (enst_dataset.py, etc.)
│   ├── utils/                  # Utility scripts (e.g., loss weight calculation)
│   ├── runs_transcription/     # Default output directory for training experiments
│   ├── GeneralUser_GS_v2_0_2/  # Soundfont for FluidSynth
│   ├── drum_transcriber_rnn.py # The model architecture definition
│   ├── train.py                # Script for training new models
│   ├── evaluate.py             # Script for quantitative model evaluation
│   ├── analyze_model.py        # Script for qualitative error analysis
│   ├── test.py                 # Script for quick inference on a single audio file
│   └── app.py                  # The Gradio interactive web application
│
├── data/                       # Directory for raw datasets (audio, MIDI, annotations)
│
├── requirements.txt            # Pip requirements file
└── conda_requirements.txt      # Conda requirements file

⚙️ Installation & Setup

1. Prerequisites

This project requires FluidSynth (version 2.4.7 or later) for the audio synthesis feature in the demo.

1. Download and install FluidSynth from the official repository.
2. Crucially, ensure the fluidsynth executable is added to your system's PATH. You should be able to run the command fluidsynth from any terminal window.

2. Clone the Repository

git clone https://github.com/GioPicci/dasheng-drum-transcribe.git
cd dasheng-drum-transcribe

3. Set Up Python Environment

You can use either Conda or a standard Python virtual environment.

Using Conda:

# Create and activate the environment
conda create --name adt_env --file conda_requirements.txt
conda activate adt_env

Using venv and pip:

# Create a virtual environment
python3 -m venv env

# Activate the environment
source env/bin/activate  # On Linux/macOS
# env\Scripts\activate   # On Windows

# Install the dependencies
pip install -r requirements.txt

4. Download Pre-trained Models

To get started quickly, you can download pre-trained models and use them directly in the demo. Four models are provided:

• GMD (Groove MIDI Dataset) - 3 Instruments
• GMD (Groove MIDI Dataset) - 7 Instruments
• ENST-Drums - 3 Instruments
• ENST-Drums - 7 Instruments

1. Download the models from this link
2. Unzip the file.
3. Copy the individual experiment folders (e.g., exp_20250807_LR5e-05_BS8_RNN384x2) into the dasheng_rnn/runs_transcription/ directory.

Once copied, they can be be automatically detected by the demo application.

5. Download the Datasets (Optional - for training)

This step is only required if you want to train your own models.

• Groove MIDI Dataset (GMD):

  • Download the dataset from the official Magenta page.
  • Extract the contents and place the folder inside the data/ directory.

• ENST-Drums Dataset:

  • This dataset is not publicly available and requires authorization for academic use.
  • You can request access by filling out this form.
  • Once you gain access, download the dataset and place it inside the data/ directory.

🚀 Usage

All scripts should be run from the dasheng_rnn/ directory.

1. Running the Interactive Demo

To launch the web application, make sure your virtual environment is active, FluidSynth is installed, and you have placed the pre-trained models in the correct folder.

python app.py

The demo will be available at http://127.0.0.1:7860.

2. Training a New Model

The train.py script handles the entire training loop. All hyperparameters can be configured via command-line arguments.

Basic Usage To start training with default settings (GMD dataset, 2-layer bidirectional GRU, etc.):

python train.py

A new experiment folder will be created in runs_transcription/, containing logs, the config.yaml, and the best model weights (.pth).

Customizing the Training Use command-line arguments to customize the experiment. To see all available options:

python train.py --help

Example 1: Train on the ENST dataset for 50 epochs with a batch size of 16.

python train.py --dataset ENST --batch_size 16 --epochs 50

Example 2: Try a different architecture (single-layer non-bidirectional LSTM) with a higher learning rate.

python train.py --rnn_type LSTM --rnn_layers 1 --rnn_hidden_size 256 --no_bidirectional --lr 1e-4

Monitoring with TensorBoard You can monitor training progress in real-time with TensorBoard:

tensorboard --logdir runs_transcription --host 0.0.0.0 --port 6006

3. Evaluating a Trained Model

The evaluate.py script provides an objective performance assessment of a trained model using a two-phase process to prevent data leakage from the test set.

Evaluation Process:

1. Phase 1: Threshold Tuning: The script first finds the optimal decision threshold for each instrument by maximizing the F1-score on the validation set. These thresholds are saved to optimal_thresholds.json.
2. Phase 2: Final Metrics: Using these optimal thresholds, the model is then evaluated on the unseen test set. It calculates Precision, Recall, and F1-score (per-instrument, macro-averaged, and micro-averaged).

How to Run Simply point the script to an experiment directory:

python evaluate.py runs_transcription/your_experiment_folder

4. Qualitative Error Analysis

While evaluate.py answers "how well" the model performs, analyze_model.py helps answer "why and where" it makes mistakes. It identifies individual False Positives and False Negatives on the validation set and logs them to error_analysis.csv.

python analyze_model.py runs_transcription/your_experiment_folder

5. Quick Test on a Single File

Use test.py for a quick inference on a single audio file without launching the full demo. The script generates a .mid file and a .png image of the piano roll in a test_results/ folder.

python test.py runs_transcription/your_experiment_folder /path/to/your/drum_track.wav