Logo
Project Icon

fMRI-reconstruction-NSD

MindEye 从fMRI数据重建和检索视觉信息

fMRI重建MindEye脑活动图像生成自然场景数据集Github开源项目

MindEye项目将fMRI数据转化为图像重建和检索结果。该项目结合对比学习和扩散先验模型,实现大脑活动到视觉信息的映射。MindEye能重建观看的图像,并从LAION数据库中检索相似图像。项目提供代码、预训练模型和使用说明,支持认知神经科学和人工智能的交叉研究。

访问官网访问官网GithubGithubHuggingfaceHuggingface论文论文
介绍相关项目

MindEye: fMRI到图像重建与检索

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


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

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

安装说明

  1. 同意自然场景数据集的条款和条件并填写NSD数据访问表格

  2. 下载此仓库: git clone https://github.com/MedARC-AI/fMRI-reconstruction-NSD.git

  3. 运行setup.sh创建一个包含运行我们脚本所需包的conda环境；使用conda activate mindeye激活环境。

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

一般信息

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

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

相关项目
项目侧边栏1项目侧边栏2
推荐项目
Project Cover

豆包MarsCode

豆包 MarsCode 是一款革命性的编程助手，通过AI技术提供代码补全、单测生成、代码解释和智能问答等功能，支持100+编程语言，与主流编辑器无缝集成，显著提升开发效率和代码质量。

Project Cover

AI写歌

Suno AI是一个革命性的AI音乐创作平台，能在短短30秒内帮助用户创作出一首完整的歌曲。无论是寻找创作灵感还是需要快速制作音乐，Suno AI都是音乐爱好者和专业人士的理想选择。

Project Cover

有言AI

有言平台提供一站式AIGC视频创作解决方案，通过智能技术简化视频制作流程。无论是企业宣传还是个人分享，有言都能帮助用户快速、轻松地制作出专业级别的视频内容。

Project Cover

Kimi

Kimi AI助手提供多语言对话支持，能够阅读和理解用户上传的文件内容，解析网页信息，并结合搜索结果为用户提供详尽的答案。无论是日常咨询还是专业问题，Kimi都能以友好、专业的方式提供帮助。

Project Cover

阿里绘蛙

绘蛙是阿里巴巴集团推出的革命性AI电商营销平台。利用尖端人工智能技术，为商家提供一键生成商品图和营销文案的服务，显著提升内容创作效率和营销效果。适用于淘宝、天猫等电商平台，让商品第一时间被种草。

Project Cover

吐司

探索Tensor.Art平台的独特AI模型，免费访问各种图像生成与AI训练工具，从Stable Diffusion等基础模型开始，轻松实现创新图像生成。体验前沿的AI技术，推动个人和企业的创新发展。

Project Cover

SubCat字幕猫

SubCat字幕猫APP是一款创新的视频播放器，它将改变您观看视频的方式！SubCat结合了先进的人工智能技术，为您提供即时视频字幕翻译，无论是本地视频还是网络流媒体，让您轻松享受各种语言的内容。

Project Cover

美间AI

美间AI创意设计平台，利用前沿AI技术，为设计师和营销人员提供一站式设计解决方案。从智能海报到3D效果图，再到文案生成，美间让创意设计更简单、更高效。

Project Cover

AIWritePaper论文写作

AIWritePaper论文写作是一站式AI论文写作辅助工具，简化了选题、文献检索至论文撰写的整个过程。通过简单设定，平台可快速生成高质量论文大纲和全文，配合图表、参考文献等一应俱全，同时提供开题报告和答辩PPT等增值服务，保障数据安全，有效提升写作效率和论文质量。

使用协议隐私政策广告服务
投诉举报邮箱: service@vectorlightyear.com
@2024 懂AI·鲁ICP备2024100362号-6·鲁公网安备37021002001498号