README.md

---
license: mit
tags:
  - audio
  - speech-enhancement
  - knowledge-distillation
  - codec2
  - lora
  - pytorch
  - edge-ai
datasets:
  - jspaulsen/vctk
metrics:
  - pesq
  - stoi
library_name: pytorch
---

# Chola-Compressor: LoRaVoiceLink Speech Restoration Model

A compact spectrogram U-Net (**1.93M parameters**) that restores speech quality lost to Codec2 compression, trained via knowledge distillation for use in off-grid LoRa voice communication links. Designed to run in real time on edge hardware (Jetson-class), not in the cloud.

## Problem

LoRa radio has enough bandwidth for Codec2 at very low bitrates (this project uses the 1200bps mode), which makes off-grid voice communication possible but introduces heavy compression artifacts. This model sits after Codec2 decoding on the receiving end and restores some of that lost quality.

## Architecture

3-level U-Net operating on STFT magnitude spectrograms (n_fft=512, hop=128, 16kHz):

- Encoder: 3 conv blocks (32→64→128 channels) with max-pooling
- Bottleneck: 256 channels
- Decoder: 3 conv blocks with transposed-conv upsampling and skip connections
- Output: predicts a `[0,1]` mask multiplied against the input magnitude (not raw magnitude directly — more stable to train)
- Phase is **not** predicted; the model reuses the degraded audio's own phase for reconstruction (see Limitations)

## Training

- **Data**: 5,000 utterances from VCTK (33 speakers, mic2 only, speaker-disjoint train/val/test split — no speaker overlap across splits), streamed from the [jspaulsen/vctk](https://huggingface.co/datasets/jspaulsen/vctk) mirror
- **Degradation**: real Codec2 encode/decode roundtrip (1200bps), not a synthetic approximation
- **Distillation**: trained with [Meta's Denoiser](https://github.com/facebookresearch/denoiser) (dns64) as an auxiliary teacher signal alongside the real clean-speech target. Ablation showed the teacher term contributed no measurable benefit over training on ground truth alone (see Results) — this checkpoint (`models_no_teacher`) was trained with `teacher_weight=0`, i.e. supervised directly against real clean speech.
- **Loss**: L1 on log-magnitude spectrograms (`log1p`), chosen after finding raw-magnitude L1 over-weights loud regions and under-penalizes quiet, perceptually important detail
- 70 epochs, Adam, lr=1e-4, batch size capped by 8GB VRAM

## Results

Evaluated on a held-out, speaker-disjoint test split (403 utterances). All three metrics computed against the true clean reference.

| Candidate | PESQ ↑ | STOI ↑ | SI-SDR (dB) ↑ |
|---|---|---|---|
| Codec2-degraded (no processing) | 1.522 | 0.658 | -28.23 |
| Denoiser (teacher, for reference) | 1.543 | 0.650 | -28.08 |
| **This model** | 1.471 | **0.780** | **-26.69** |

![Benchmark results](chola_compressor_benchmark_results.png)

![Spectrogram comparison from a live voice test](chola_compressor_spectrogram_comparison.png)

**Real, verified gains:** +0.12 STOI (intelligibility) and +1.5dB SI-SDR over doing nothing. Both metrics are dominated by energy/envelope accuracy, where this model clearly helps.

**Known limitation — PESQ:** PESQ is highly sensitive to phase accuracy, and this model only predicts a magnitude mask, reconstructing with the degraded audio's original (uncontrolled) phase. That's the most likely explanation for PESQ landing slightly below the unprocessed baseline despite STOI/SI-SDR improving substantially — three independent loss-function reformulations (distillation weight, log vs. linear magnitude) all left PESQ in the same 1.47–1.48 range, which is consistent with a phase-reconstruction ceiling rather than a loss-tuning problem. A phase-aware architecture (predicting complex spectrograms or a phase correction term) would likely be needed to close this gap; that's a known next step, not yet implemented in this checkpoint.

![PESQ stayed flat across three loss-function ablations](chola_compressor_pesq_ablation.png)

## Intended use

Research and portfolio demonstration of magnitude-domain speech restoration via knowledge distillation for bandwidth-constrained voice links. Not validated for safety-critical or emergency-communication deployment.

## How to use

```python
import torch
import numpy as np
import librosa
import soundfile as sf

N_FFT, HOP_LENGTH, SAMPLE_RATE = 512, 128, 16000

class SpectrogramUNet(torch.nn.Module):
    # ... see model.py in the project repo for the full class definition
    pass

def restore(degraded_wav_path, model, device="cpu"):
    audio, sr = sf.read(degraded_wav_path, dtype="float32")
    if sr != SAMPLE_RATE:
        audio = librosa.resample(audio, orig_sr=sr, target_sr=SAMPLE_RATE)
    stft = librosa.stft(audio, n_fft=N_FFT, hop_length=HOP_LENGTH)
    mag, phase = np.abs(stft), np.angle(stft)

    x = torch.from_numpy(mag).float().unsqueeze(0).unsqueeze(0).to(device)
    with torch.no_grad():
        pred_mag = model(x).squeeze().cpu().numpy()

    restored_stft = pred_mag * np.exp(1j * phase)  # reuses degraded audio's phase
    return librosa.istft(restored_stft, hop_length=HOP_LENGTH)

model = SpectrogramUNet(base_channels=32)
model.load_state_dict(torch.load("best.pt", map_location="cpu"))
model.eval()

restored = restore("degraded.wav", model)
sf.write("restored.wav", restored, SAMPLE_RATE)
```

Full training/eval/live-test code: see the project repository.

## Citation

If you use this model, please cite the LoRaVoiceLink project (link to source repo).