CLARITY-Net: A Lightweight Causal U-Net for Real-Time Speech Enhancement

CLARITY-Net is a speech denoising project built around a lightweight causal U-Net that works on short-time Fourier transform (STFT) spectrograms. The repository combines the research paper in CLARITY Net.pdf with the full notebook-based pipeline for dataset preparation, model training, inference, and evaluation.

The goal of the project is to improve speech quality in noisy environments while keeping latency and model size low enough for real-time or edge-device deployment.

Overview

The codebase implements a complete speech enhancement workflow:

1. Prepare clean speech clips and generate noisy mixtures.
2. Train a causal U-Net-style enhancement model on log-magnitude STFT features.
3. Run denoising on the test set and save enhanced waveforms.
4. Evaluate the output with objective speech-quality metrics.

The research paper reports the following headline results for the proposed model:

• 63,896 trainable parameters
• SI-SDR: 13.7 dB
• PESQ: 2.865 (narrowband) and 2.135 (wideband)
• STOI: 0.906
• Real-time factor (RTF): 0.309

Repository Layout

.
├── CLARITY Net.pdf
├── architecture.ipynb
├── evaluation.ipynb
├── generate.ipynb
├── requirements.txt
└── dataset/
    ├── clean/
    ├── clean_clipped/
    ├── noisy/
    ├── noisy_pink/
    ├── noisy_white/
    ├── noisy-2/
    └── test/
        ├── clean/
        ├── enhanced/
        └── noisy/

Notebook Roles

• generate.ipynb prepares the training and test mixtures.
• architecture.ipynb defines the CLARITY-Net model and trains it.
• evaluation.ipynb benchmarks denoising quality and computes speech metrics.

Method Summary

The model operates in the time-frequency domain:

• Audio is resampled to 16 kHz.
• STFT is applied with Hann windows.
• The network consumes log-magnitude spectrograms.
• The model predicts a residual correction to the noisy magnitude.
• The enhanced waveform is reconstructed with the original noisy phase.

The paper describes CLARITY-Net as a causal U-Net with:

• depthwise separable convolutions,
• squeeze-and-excitation channel attention,
• group normalization,
• a bottleneck dilation block for a larger receptive field,
• residual learning with a scaling factor of 0.3.

This design keeps the model small while preserving temporal causality for streaming use.

Paper-Derived Model Description

CLARITY-Net is described in the paper as a lightweight causal U-Net for real-time speech enhancement. The model learns to map a noisy signal $y(t) = x(t) + n(t)$ back toward the clean target $x(t)$, using STFT features instead of raw waveform regression.

Input Representation

The paper uses a Hann-windowed STFT with:

• FFT size: 512
• hop length: 256 samples
• window length: 512 samples
• sample rate: 16 kHz

This yields a spectrogram patch shaped like 257 x 24 in the default notebook implementation, where frequency bins are processed over a short causal time window.

Residual Magnitude Learning

Rather than predicting the clean magnitude directly, the model predicts a residual correction $R(f, \tau)$ over the noisy magnitude. The enhanced magnitude is computed as:

$$ |\hat{X}(f, \tau)| = \max(|Y(f, \tau)| + \alpha R(f, \tau), \epsilon) $$

with the paper using $\alpha = 0.3$ and a small $\epsilon$ for numerical stability.

The noisy phase is reused during reconstruction. This keeps the system simple and consistent with the paper’s real-time design goal.

Causal U-Net Backbone

The network is a compact encoder-decoder with skip connections:

• encoder stages downsample only along the frequency axis,
• the time axis remains causal,
• a bottleneck expands the receptive field using dilation,
• the decoder upsamples the frequency axis and fuses skip features,
• the final output is a residual spectrogram estimate.

The paper’s reference configuration is:

• input: 257 x 24
• encoder channels: [12, 24, 48]
• bottleneck channels: 96
• output: 257 x 24

Causal and Efficient Convolutions

To support streaming inference, the paper uses causal padding on the time axis so each prediction depends only on current and past context. To reduce compute, the convolutions are depthwise separable:

• depthwise convolution handles per-channel spatial filtering,
• pointwise convolution mixes channel information,
• GroupNorm is used for stability with small batch sizes.

Channel Attention

The model includes Squeeze-and-Excitation blocks that apply channel-wise attention. These blocks compress each feature map with global average pooling, then learn a channel scale through a small bottleneck MLP before reweighting the activations.

This gives the network a lightweight way to emphasize informative channels without adding much cost.

Training Objective

The paper presents a hybrid objective that combines magnitude-domain reconstruction with SI-SDR-based waveform quality. In the notebook implementation, the optimization is kept compact and stable by emphasizing an L1 loss on the predicted magnitude residuals, while the paper describes a broader composite training strategy.

Reference Diagrams

The paper includes the following core figures, which are useful when reading the implementation:

• Fig. 1: High-level CLARITY-Net architecture showing the encoder, bottleneck, decoder, and skip connections.
• Fig. 2: The Squeeze-and-Excitation block used for channel attention.
• Fig. 3: The convolutional block used throughout the network.

The notebook implementation follows the same structure, so the figures in the paper map directly to the model classes in architecture.ipynb.

Pipeline Diagram

flowchart LR
    A[Clean speech] --> B[Add ESC-10 / white / pink noise]
    B --> C[STFT magnitude + phase]
    C --> D[Log-magnitude spectrogram]
    D --> E[CLARITY-Net causal U-Net]
    E --> F[Residual magnitude estimate]
    F --> G[Reuse noisy phase]
    G --> H[iSTFT reconstruction]
    H --> I[Enhanced speech]

Loading

Architecture Diagram

flowchart TB
    I[Input 257 x 24] --> E1[Encoder 1 12 ch]
    E1 --> E2[Encoder 2 24 ch]
    E2 --> E3[Encoder 3 48 ch]
    E3 --> B[Bottleneck 96 ch + dilation]
    B --> D3[Decoder 3 48 ch]
    D3 --> D2[Decoder 2 24 ch]
    D2 --> D1[Decoder 1 12 ch]
    D1 --> O[Output residual 257 x 24]

Loading

Dataset Preparation

The project uses a composite dataset built from:

• clean speech from LibriSpeech,
• environmental noise from ESC-10,
• synthetic white noise,
• synthetic pink noise.

The generation pipeline in generate.ipynb produces noisy mixtures at multiple SNR levels and organizes them into the dataset folders used by training and evaluation.

Common dataset folders used by the notebooks are:

• dataset/clean/ - original clean speech
• dataset/clean_clipped/ - clean clips trimmed to a common length
• dataset/noisy_white/ - white-noise mixtures
• dataset/noisy_pink/ - pink-noise mixtures
• dataset/noisy/ - final training noisy set
• dataset/test/noisy/ - evaluation noisy set
• dataset/test/clean/ - matching clean references for the test set
• dataset/test/enhanced/ - model outputs written during inference

Environment Setup

Install the Python dependencies with:

pip install -r requirements.txt

The project was developed in a Python 3.11 virtual environment. If you want to reproduce the notebook workflow locally, make sure Jupyter support is available in the same environment.

How to Run

1. Generate the noisy dataset

Open generate.ipynb and run the cells in order. This notebook:

• measures the shortest clean clip length,
• clips all clean files to a common duration,
• mixes in white and pink noise at multiple SNR levels,
• builds the ESC-10 based noisy set,
• creates balanced train/test splits under dataset/noisy/ and dataset/test/.

2. Train CLARITY-Net

Open architecture.ipynb and run the training cells. This notebook:

• validates the dataset folders,
• builds the AudioDataset class,
• converts audio into STFT magnitude and phase pairs,
• defines the lightweight causal U-Net,
• trains the model and saves the best checkpoint as best_denoiser.pth,
• exports a portable model payload as best_denoiser.pkl.

The training notebook uses residual learning on magnitude spectrograms and reconstructs audio with the noisy phase.

3. Run inference on the test set

The inference section in architecture.ipynb loads the best checkpoint and writes enhanced .wav files to dataset/test/enhanced/.

It also includes a streaming-style latency benchmark that estimates average per-chunk latency and real-time factor.

4. Evaluate enhancement quality

Open evaluation.ipynb to compute objective metrics on the test set. The notebook reports:

• PESQ
• STOI
• SI-SDR
• SNR / SI-SNR
• composite measures such as CSIG, CBAK, and COVL when available

Expected Outputs

After running the notebooks, you should see:

• a trained model checkpoint (best_denoiser.pth),
• a serialized model export (best_denoiser.pkl),
• denoised audio in dataset/test/enhanced/,
• metric tables and benchmark prints inside the evaluation notebook.

Architecture Notes

The implemented model is intentionally compact:

• input shape: 257 x 24 spectrogram patches,
• encoder: 3 layers with channels [12, 24, 48],
• bottleneck: 96 channels with a dilated block,
• decoder: symmetric upsampling path,
• output: residual spectrogram estimate.

The paper emphasizes causal processing and low latency, making the model suitable for streaming-style enhancement rather than offline, full-context denoising.

Evaluation Notes

The paper reports that CLARITY-Net was compared with stronger baseline families, including:

• classical enhancement methods,
• vanilla U-Net,
• CRN / DCCRN-style models,
• Conv-TasNet,
• other lightweight speech enhancement models.

The intent of the evaluation is to balance speech quality against computational cost.

Paper References

The paper also frames the method around the following ideas:

• real-time speech enhancement under latency constraints,
• causal processing for streaming deployment,
• efficiency gains from depthwise separable convolutions,
• feature recalibration using SE attention,
• reconstruction quality evaluated with SI-SDR, PESQ, and STOI.

These ideas are the best guide for reading both the notebook code and the included PDF.

Reproducibility Tips

• Keep the dataset folder structure unchanged.
• Make sure training and test file naming stays consistent between noisy and clean pairs.
• If you regenerate the dataset, rerun the notebooks in the order: generate.ipynb, architecture.ipynb, evaluation.ipynb.
• On Apple Silicon, the notebooks prefer mps when available and fall back to CPU when needed.

References

• CLARITY Net.pdf - the research paper for this implementation.
• architecture.ipynb - model architecture, training, and inference.
• generate.ipynb - dataset creation and noisy mixture synthesis.
• evaluation.ipynb - benchmark and metric computation.

License

No explicit license file is present in the repository. Add one if you want to publish or share the project.