MindEye: fMRI到图像重建与检索

2024年5月22日: 查看我们的新作品 MindEye2，在所有指标上都超越了MindEye1!

<br>

项目页面: https://medarc-ai.github.io/mindeye/

arXiv预印本: https://arxiv.org/abs/2305.18274

安装说明

同意自然场景数据集的条款和条件并填写NSD数据访问表格
下载此仓库: git clone https://github.com/MedARC-AI/fMRI-reconstruction-NSD.git
运行setup.sh创建一个包含运行我们脚本所需包的conda环境；使用conda activate mindeye激活环境。

cd fMRI-reconstruction-NSD/src
. setup.sh

(可选) 对于LAION-5B检索，你还需要额外下载预训练的检查点。为此，进入"src"文件夹并运行. download.sh。这将允许你从预训练检查点（从LAION-Aesthetics训练的文本到图像扩散先验）开始训练扩散先验。我们观察到，使用这个检查点，而不是从头开始训练先验，显著改善了LAION-5B检索。

一般信息

此仓库包含用于以下目的的Jupyter笔记本：

训练MindEye (src/Train_MindEye.ipynb)
使用训练好的模型从大脑活动重建图像 (src/Reconstructions.ipynb)
从大脑活动中检索图像，可以是测试集或通过LAION-5B (src/Retrievals.ipynb)
根据低级和高级图像指标评估重建效果与真实图像的对比 (src/Reconstruction_Metrics.ipynb)

以上所有Jupyter笔记本也有相应的python (.py) 文件，可以通过命令行运行。

这个仓库还包含将大脑活动映射到Stable Diffusion变分自编码器的代码 (src/train_autoencoder.py)。

受试者1的预训练模型

你可以跳过自己训练MindEye，而是在NSD的受试者1上运行其余的笔记本，方法是下载我们在huggingface上提供的预训练模型，并将这些包含模型检查点的文件夹放在"fMRI-reconstruction-NSD/train_logs/"内。

prior_257_final_subj01_bimixco_softclip_byol: CLIP ViT-L/14隐藏层 (257x768) 
prior_1x768_final_subj01_bimixco_softclip_byol: CLIP ViT-L/14最终层 (1x768)
autoencoder_subj01_4x_locont_no_reconst: Stable Diffusion VAE (低级管道)

训练MindEye (高级管道)

通过Train_MindEye.py训练MindEye。

将data_path设置为包含自然场景数据集的文件夹（如果未找到将下载到那里；每个受试者>30Gb，仅下载当前受试者的数据）。
将model_name设置为你想给模型命名的名称，用于保存。
如果你想映射到CLIP的最终层以进行LAION-5B检索或通过Stable Diffusion（图像变体）进行重建，请设置--no-hidden --no-norm_embs。否则，Versatile Diffusion使用默认的--hidden --norm_embs。

可以设置各种参数进行训练（见下文）；默认是将MindEye训练到CLIP ViT-L/14的最后一个隐藏层，使用与我们论文相同的设置，针对NSD的受试者1。

训练好的模型检查点将保存在"fMRI-reconstruction-NSD/train_logs"文件夹内。所有其他输出都保存在"fMRI-reconstruction-NSD/src"文件夹内。

$ python Train_MindEye.py --help

usage: Train_MindEye.py [-h] [--model_name MODEL_NAME] [--data_path DATA_PATH]
                        [--subj {1,2,5,7}] [--batch_size BATCH_SIZE]
                        [--hidden | --no-hidden]
                        [--clip_variant {RN50,ViT-L/14,ViT-B/32,RN50x64}]
                        [--wandb_log | --no-wandb_log]
                        [--resume_from_ckpt | --no-resume_from_ckpt]
                        [--wandb_project WANDB_PROJECT]
                        [--mixup_pct MIXUP_PCT] [--norm_embs | --no-norm_embs]
                        [--use_image_aug | --no-use_image_aug]
                        [--num_epochs NUM_EPOCHS] [--prior | --no-prior]
                        [--v2c | --no-v2c] [--plot_umap | --no-plot_umap]
                        [--lr_scheduler_type {cycle,linear}]
                        [--ckpt_saving | --no-ckpt_saving]
                        [--ckpt_interval CKPT_INTERVAL]
                        [--save_at_end | --no-save_at_end] [--seed SEED]
                        [--max_lr MAX_LR] [--n_samples_save {0,1}]
                        [--use_projector | --no-use_projector]
                        [--vd_cache_dir VD_CACHE_DIR]

模型训练配置

选项:
  -h, --help            显示此帮助消息并退出
  --model_name MODEL_NAME
                        模型名称，用于ckpt保存和wandb日志记录（如果启用）
  --data_path DATA_PATH
                        NSD数据存储/下载的路径
  --subj {1,2,5,7}
  --batch_size BATCH_SIZE
                        如果只训练v2c而不训练扩散先验，批量大小可以增加10倍
  --hidden, --no-hidden
                        如果为True，CLIP嵌入将来自最后一个隐藏层（例如，257x768 - Versatile Diffusion），
                        而不是最终层（默认：True）
  --clip_variant {RN50,ViT-L/14,ViT-B/32,RN50x64}
                        OpenAI clip变体
  --wandb_log, --no-wandb_log
                        是否记录到wandb（默认：False）
  --resume_from_ckpt, --no-resume_from_ckpt
                        如果不使用wandb并想从ckpt恢复（默认：False）
  --wandb_project WANDB_PROJECT
                        wandb项目名称
  --mixup_pct MIXUP_PCT
                        训练过程中从BiMixCo切换到SoftCLIP的比例
  --norm_embs, --no-norm_embs
                        对CLIP嵌入进行l2归一化（默认：True）
  --use_image_aug, --no-use_image_aug
                        是否使用图像增强（默认：True）
  --num_epochs NUM_EPOCHS
                        训练的轮数
  --prior, --no-prior   如果为False，将仅使用CLIP损失并忽略扩散先验（默认：True）
  --v2c, --no-v2c       如果为False，将仅使用扩散先验损失（默认：True）
  --plot_umap, --no-plot_umap
                        绘制UMAP图以及重建（默认：False）
  --lr_scheduler_type {cycle,linear}
  --ckpt_saving, --no-ckpt_saving
  --ckpt_interval CKPT_INTERVAL
                        每x个epoch保存备份ckpt并重建
  --save_at_end, --no-save_at_end
                        如果为True，在训练结束时保存best.ckpt。如果为False且ckpt_saving==True，
                        将在epoch显示最佳验证分数时保存best.ckpt（默认：False）
  --seed SEED
  --max_lr MAX_LR
  --n_samples_save {0,1}
                        用于监控进度的重建数量，0将加快训练速度
  --use_projector, --no-use_projector
                        在主MLP之后添加额外的MLP，以便模型可以单独学习一种方法来最小化先验损失（BYOL）的NCE（默认：True）
  --vd_cache_dir VD_CACHE_DIR
                        缓存的Versatile Diffusion模型位置；如果未缓存将下载到此路径

从预训练的MindEye重建

现在你在"train_logs"文件夹中有了预训练的模型ckpts，无论是通过运行Train_MindEye.py还是通过从huggingface下载我们预训练的受试者1模型，我们都可以继续从测试集的未见过的大脑活动中重建图像。

Reconstructions.py默认将Versatile Diffusion重建输出为torch .pt文件，不使用img2img，也不使用二阶选择（recons_per_sample=1）。

将data_path设置为包含自然场景数据集的文件夹（需要先下载，可以通过上述训练脚本或手动从huggingface下载）。
将model_name设置为"fMRI-reconstruction-NSD/train_logs"中包含将大脑活动映射到CLIP最后隐藏层的ckpt的文件夹名称。
如果你想使用img2img，将autoencoder_name设置为"fMRI-reconstruction-NSD/train_logs"中包含将大脑活动映射到Stable Diffusion变分自编码器的ckpt的文件夹名称。
如果你使用img2img，将img2img_strength设置为你喜欢的指导级别，其中1=无img2img，0=仅输出低级管道的结果。

$ python Reconstructions.py --help

usage: Reconstructions.py [-h] [--model_name MODEL_NAME]
                          [--autoencoder_name AUTOENCODER_NAME] [--data_path DATA_PATH]
                          [--subj {1,2,5,7}] [--img2img_strength IMG2IMG_STRENGTH]
                          [--recons_per_sample RECONS_PER_SAMPLE]
                          [--vd_cache_dir VD_CACHE_DIR]

模型训练配置

选项:
  -h, --help            显示此帮助消息并退出
  --model_name MODEL_NAME
                        训练好的模型名称
  --autoencoder_name AUTOENCODER_NAME
                        训练好的自编码器模型名称
  --data_path DATA_PATH
                        NSD数据存储的路径（见README）
  --subj {1,2,5,7}
  --img2img_strength IMG2IMG_STRENGTH
                        img2img的强度（1=无img2img；0=输出低级图像本身）
  --recons_per_sample RECONS_PER_SAMPLE
                        输出多少个重建，然后自动选择最佳的一个（MindEye使用16）
  --vd_cache_dir VD_CACHE_DIR
                        缓存的Versatile Diffusion模型位置；如果未缓存将下载到此路径

图像/大脑检索（包括LAION-5B图像检索）

要使用NSD测试集评估图像/大脑检索，请使用Jupyter笔记本Retrievals.ipynb并按照"图像/大脑检索"标题下的代码块进行操作。

运行Retrievals.py将基于MindEye变体（将大脑活动映射到CLIP的最终层）检索LAION-5B中的前16个最近邻。接着进行二阶选择，将16个检索到的图像转换为CLIP最后隐藏层嵌入，并与核心模型（将大脑活动映射到CLIP最后隐藏层）的MindEye输出进行比较。将选择CLIP相似度最高的检索图像，所有top-1检索结果将保存为torch .pt文件。

将 data_path 设置为包含自然场景数据集的文件夹（如果未找到将下载到该处；每个受试者超过30GB，仅下载当前受试者的数据）。
将 model_name 设置为 "fMRI-reconstruction-NSD/train_logs" 中包含将大脑活动映射到 CLIP 最后一个隐藏层的 ckpt 文件的文件夹名称。
将 model_name2 设置为 "fMRI-reconstruction-NSD/train_logs" 中包含将大脑活动映射到 CLIP 最终层的 ckpt 文件的文件夹名称。

$ python Retrievals.py --help

用法: Retrievals.py [-h] [--model_name MODEL_NAME]
                               [--model_name2 MODEL_NAME2] [--data_path DATA_PATH]
                               [--subj {1,2,5,7}]

模型训练配置

选项:
  -h, --help            显示此帮助信息并退出
  --model_name MODEL_NAME
                        257x768 模型的名称，用于除 LAION-5B 检索之外的所有操作
  --model_name2 MODEL_NAME2
                        1x768 模型的名称，用于 LAION-5B 检索
  --data_path DATA_PATH
                        NSD 数据存储的路径（参见 README）
  --subj {1,2,5,7}

评估重建结果

在运行 Reconstructions.py 或 Retrievals.py 保存 .pt 文件后，您可以使用 Reconstruction_Metrics.py 来评估重建图像，使用与论文中相同的低级和高级图像指标。

将 recon_path 设置为 "fMRI-reconstruction-NSD/src" 中由 Reconstructions.py 输出的文件名（应为 {model_name}_recons_img2img{img2img_strength}_{recons_per_sample}samples.pt）。
或者，要评估 LAION-5B 检索结果，您可以用 Retrievals.py 输出的 .pt 文件名替换 recon_path（应为 {model_name}_laion_retrievals_top16.pt）。
将 all_images_path 设置为 "fMRI-reconstruction-NSD/src" 中由 Reconstructions.py 或 Retrievals.py 输出的 all_images.pt 文件（应为 all_images.pt）。

$ python Reconstruction_Metrics.py --help

用法: Reconstruction_Metrics.py [-h] [--recon_path RECON_PATH]
                                 [--all_images_path ALL_IMAGES_PATH]

模型训练配置

选项:
  -h, --help            显示此帮助信息并退出
  --recon_path RECON_PATH
                        重建/检索输出的路径
  --all_images_path ALL_IMAGES_PATH
                        真实图像输出的路径

训练 MindEye（低级流水线）

运行 train_autoencoder.py 来训练 MindEye 低级模型。

在训练之前，请在 Python 文件中将 train_url、val_url 和 meta_url 变量设置为相关的 NSD 数据集位置。训练代码需要 VICRegL ConvNext-XL 的权重（从这里下载）和 Stable Diffusion Image Variations 自编码器的权重（从这里下载）。请注意，这些 SD 自编码器权重是从 LambdaLabs SD Image Variations v2 模型这里原样提取的。

引用

如果您使用了这项工作，请同时引用 MindEye 论文和自然场景数据集论文。同时也建议引用我们最新的 MindEye2 论文，它对 MindEye1 进行了改进。

MindEye2

Scotti, Tripathy, Torrico, Kneeland, Chen, Narang, Santhirasegaran, Xu, Naselaris, Norman, & Abraham. MindEye2: Shared-Subject Models Enable fMRI-To-Image With 1 Hour of Data. International Conference on Machine Learning. (2024). arXiv:2403.11207

MindEye1

Scotti, Banerjee, Goode, Shabalin, Nguyen, Cohen, Dempster, Verlinde, Yundler, Weisberg, Norman, & Abraham. Reconstructing the Mind's Eye: fMRI-to-Image with Contrastive Learning and Diffusion Priors. Advances in Neural Information Processing Systems, 36. (2023). arXiv:2305.18274.

自然场景数据集

Allen, St-Yves, Wu, Breedlove, Prince, Dowdle, Nau, Caron, Pestilli, Charest, Hutchinson, Naselaris, & Kay. A massive 7T fMRI dataset to bridge cognitive neuroscience and artificial intelligence. Nature Neuroscience (2021).