/BrainCodec

For paper

Primary LanguagePythonMIT LicenseMIT

BrainCraft

BrainCodec Architecture

BrainCraft is a PyTorch library based on AudioCraft. BrainCraft currently includes the inference and training code for BrainCodec.

Datasets

  • For upstream and downstream tasks:
  • For resting-state fMRI:
    • Download from here.
    • Note: Due to licensing restrictions, the SRPBS Traveling Subject MRI Dataset mentioned in the paper cannot be redistributed. You need to download it yourself and create a tarfile.

Models

At the moment, BrainCraft contains the training code and inference code for:

  • BrainCodec (ours): A state-of-the-art neural codec model for fMRI data
  • CSM (Thomas et al. (2022)) : A baseline model for decoding of cognitive brain states
  • CSM+BrainCodec (ours): A state-of-the-art decoding model of cognitive brain states

Additionally, the list of available pre-trained weights is as follows:

Model Downstream Pre-trained Weight
BrainCodec - Link
CSM HCP Link
CSM MDTB Link
CSM Over100 Link
CSM + BrainCodec HCP Link
CSM + BrainCodec MDTB Link
CSM + BrainCodec Over100 Link

Installation

BrainCraft requires Python 3.9 and PyTorch 2.1.0. To install BrainCraft, you can run the following:

# Best to make sure you have torch installed first, in particular before installing xformers.
# Don't run this if you already have PyTorch installed.
pip install 'torch>=2.1'
pip install -e .

We also recommend having ffmpeg installed, either through your system or Anaconda:

sudo apt-get install ffmpeg
# Or if you are using Anaconda or Miniconda
conda install "ffmpeg<5" -c conda-forge

License

Citation

For the general framework of BrainCraft, please cite the following.

Not yet published

Acknowledgement

This repository is implemented based on AudioCraft.