fMRI-reconstruction-NSD

fMRI-reconstruction-NSD

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

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

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

MindEye: fMRI到图像重建与检索

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

<br>

项目页面: 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.pyRetrievals.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.pyRetrievals.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_urlval_urlmeta_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).

编辑推荐精选

讯飞智文

讯飞智文

一键生成PPT和Word,让学习生活更轻松

讯飞智文是一个利用 AI 技术的项目,能够帮助用户生成 PPT 以及各类文档。无论是商业领域的市场分析报告、年度目标制定,还是学生群体的职业生涯规划、实习避坑指南,亦或是活动策划、旅游攻略等内容,它都能提供支持,帮助用户精准表达,轻松呈现各种信息。

AI办公办公工具AI工具讯飞智文AI在线生成PPTAI撰写助手多语种文档生成AI自动配图热门
讯飞星火

讯飞星火

深度推理能力全新升级,全面对标OpenAI o1

科大讯飞的星火大模型,支持语言理解、知识问答和文本创作等多功能,适用于多种文件和业务场景,提升办公和日常生活的效率。讯飞星火是一个提供丰富智能服务的平台,涵盖科技资讯、图像创作、写作辅助、编程解答、科研文献解读等功能,能为不同需求的用户提供便捷高效的帮助,助力用户轻松获取信息、解决问题,满足多样化使用场景。

热门AI开发模型训练AI工具讯飞星火大模型智能问答内容创作多语种支持智慧生活
Spark-TTS

Spark-TTS

一种基于大语言模型的高效单流解耦语音令牌文本到语音合成模型

Spark-TTS 是一个基于 PyTorch 的开源文本到语音合成项目,由多个知名机构联合参与。该项目提供了高效的 LLM(大语言模型)驱动的语音合成方案,支持语音克隆和语音创建功能,可通过命令行界面(CLI)和 Web UI 两种方式使用。用户可以根据需求调整语音的性别、音高、速度等参数,生成高质量的语音。该项目适用于多种场景,如有声读物制作、智能语音助手开发等。

Trae

Trae

字节跳动发布的AI编程神器IDE

Trae是一种自适应的集成开发环境(IDE),通过自动化和多元协作改变开发流程。利用Trae,团队能够更快速、精确地编写和部署代码,从而提高编程效率和项目交付速度。Trae具备上下文感知和代码自动完成功能,是提升开发效率的理想工具。

AI工具TraeAI IDE协作生产力转型热门
咔片PPT

咔片PPT

AI助力,做PPT更简单!

咔片是一款轻量化在线演示设计工具,借助 AI 技术,实现从内容生成到智能设计的一站式 PPT 制作服务。支持多种文档格式导入生成 PPT,提供海量模板、智能美化、素材替换等功能,适用于销售、教师、学生等各类人群,能高效制作出高品质 PPT,满足不同场景演示需求。

讯飞绘文

讯飞绘文

选题、配图、成文,一站式创作,让内容运营更高效

讯飞绘文,一个AI集成平台,支持写作、选题、配图、排版和发布。高效生成适用于各类媒体的定制内容,加速品牌传播,提升内容营销效果。

热门AI辅助写作AI工具讯飞绘文内容运营AI创作个性化文章多平台分发AI助手
材料星

材料星

专业的AI公文写作平台,公文写作神器

AI 材料星,专业的 AI 公文写作辅助平台,为体制内工作人员提供高效的公文写作解决方案。拥有海量公文文库、9 大核心 AI 功能,支持 30 + 文稿类型生成,助力快速完成领导讲话、工作总结、述职报告等材料,提升办公效率,是体制打工人的得力写作神器。

openai-agents-python

openai-agents-python

OpenAI Agents SDK,助力开发者便捷使用 OpenAI 相关功能。

openai-agents-python 是 OpenAI 推出的一款强大 Python SDK,它为开发者提供了与 OpenAI 模型交互的高效工具,支持工具调用、结果处理、追踪等功能,涵盖多种应用场景,如研究助手、财务研究等,能显著提升开发效率,让开发者更轻松地利用 OpenAI 的技术优势。

Hunyuan3D-2

Hunyuan3D-2

高分辨率纹理 3D 资产生成

Hunyuan3D-2 是腾讯开发的用于 3D 资产生成的强大工具,支持从文本描述、单张图片或多视角图片生成 3D 模型,具备快速形状生成能力,可生成带纹理的高质量 3D 模型,适用于多个领域,为 3D 创作提供了高效解决方案。

3FS

3FS

一个具备存储、管理和客户端操作等多种功能的分布式文件系统相关项目。

3FS 是一个功能强大的分布式文件系统项目,涵盖了存储引擎、元数据管理、客户端工具等多个模块。它支持多种文件操作,如创建文件和目录、设置布局等,同时具备高效的事件循环、节点选择和协程池管理等特性。适用于需要大规模数据存储和管理的场景,能够提高系统的性能和可靠性,是分布式存储领域的优质解决方案。

下拉加载更多