matijama/bird-mixit-pytorch

Bird MixIT (PyTorch)

An unofficial PyTorch port of bird MixIT, Google's sound separation model for bird recordings. It splits a recording into 4 or 8 sources, so that overlapping birds and background noise can be listened to, or classified, separately. The weights are a direct conversion of Google's official TensorFlow checkpoints (Apache 2.0), and the PyTorch model reproduces the TensorFlow outputs to float round-off.

Code, examples and documentation: https://github.com/matijama/bird-mixit-pytorch

Usage

pip install git+https://github.com/matijama/bird-mixit-pytorch

from bird_mixit_pytorch import BirdMixIT, load_audio, save_audio

model = BirdMixIT.from_pretrained(num_sources=4)        # downloads these weights once, then cached
audio = load_audio("recording.wav")                     # any sample rate -> mono at 22.05 kHz
sources = model(audio)                                  # [1, 4, samples]; the sources sum to the input

for i, source in enumerate(sources[0]):
    save_audio(f"source{i}.wav", source)

BirdMixIT.from_pretrained(num_sources=8) loads the 8-source model. Run the model in fp32: it is sensitive to reduced precision, and fp16 changes the separated sources noticeably. For speed on a GPU, use torch.compile(model).

The GitHub repository also has a fine-tuning example that adapts the model to your own recordings with mixture invariant training, which needs no reference sources.

Files

file	content
bird_mixit_4sources.safetensors	model with 4 output sources (37 MB, fp32)
bird_mixit_8sources.safetensors	model with 8 output sources (38 MB, fp32)

Citation

@inproceedings{denton2022improving,
  title     = {Improving Bird Classification with Unsupervised Sound Separation},
  author    = {Denton, Tom and Wisdom, Scott and Hershey, John R.},
  booktitle = {IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
  year      = {2022}
}

@inproceedings{wisdom2020unsupervised,
  title     = {Unsupervised Sound Separation Using Mixture Invariant Training},
  author    = {Wisdom, Scott and Tzinis, Efthymios and Erdogan, Hakan and Weiss, Ron J. and Wilson, Kevin and Hershey, John R.},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2020}
}

Licence

Apache License 2.0, the same as the original checkpoints. The weights are a format conversion of Google's release: the values are unchanged, and only the storage format and tensor layout differ.