MindEye 从fMRI数据重建和检索视觉信息
MindEye项目将fMRI数据转化为图像重建和检索结果。该项目结合对比学习和扩散先验模型,实现大脑活动到视觉信息的映射。MindEye能重建观看的图像,并从LAION数据库中检索相似图像。项目提供代码、预训练模型和使用说明,支持认知神经科学和人工智能的交叉研究。
2024年5月22日: 查看我们的新作品 MindEye2,在所有指标上都超越了MindEye1!
<br>
项目页面: https://medarc-ai.github.io/mindeye/
arXiv预印本: https://arxiv.org/abs/2305.18274
下载此仓库: git clone https://github.com/MedARC-AI/fMRI-reconstruction-NSD.git
运行setup.sh
创建一个包含运行我们脚本所需包的conda环境;使用conda activate mindeye
激活环境。
cd fMRI-reconstruction-NSD/src . setup.sh
. download.sh
。这将允许你从预训练检查点(从LAION-Aesthetics训练的文本到图像扩散先验)开始训练扩散先验。我们观察到,使用这个检查点,而不是从头开始训练先验,显著改善了LAION-5B检索。此仓库包含用于以下目的的Jupyter笔记本:
以上所有Jupyter笔记本也有相应的python (.py) 文件,可以通过命令行运行。
这个仓库还包含将大脑活动映射到Stable Diffusion变分自编码器的代码 (src/train_autoencoder.py)。
你可以跳过自己训练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 (低级 管道)
通过Train_MindEye.py
训练MindEye。
data_path
设置为包含自然场景数据集的文件夹(如果未找到将下载到那里;每个受试者>30Gb,仅下载当前受试者的数据)。model_name
设置为你想给模型命名的名称,用于保存。--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模型位置;如果未缓存将下载到此路径
现在你在"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的文件夹名称。autoencoder_name
设置为"fMRI-reconstruction-NSD/train_logs"中包含将大脑活动映射到Stable Diffusion变分自编码器的ckpt的文件夹名称。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模型位置;如果未缓存将下载到此路径
要使用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
)。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
真实图像输出的路径
运行 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).
一键生成PPT和Word,让学习生活更轻松
讯飞智文是一个利用 AI 技术的项目,能够帮助用户生成 PPT 以及各类文档。无论是商业领域的市场分析报告、年度目标制定,还是学生群体的职业生涯规划、实习避坑指南,亦或是活动策划、旅游攻略等内容,它都能提供支持,帮助用户精准表达,轻松呈现各种信息。
深度推理能力全新升级,全面对标OpenAI o1
科大讯飞的星火大模型,支持语言理解、知识问答和文本创作等多功能,适用于多种文件和业务场景,提升办公和日常生活的效率。讯飞星火是一个提供丰富智能服务的平台,涵盖科技资讯、图像创作、写作辅助、编程解答、科研文献解读等功能,能为不同需求的用户提供便捷高效的帮助,助力用户轻松获取信息、解决问题,满足多样化使用场景。
一种基于大语言模型的高效单流解耦语音令牌文本到语音合成模型
Spark-TTS 是一个基于 PyTorch 的开源文本到语音合成项目,由多个知名机构联合参与。该项目提供了高效的 LLM(大语言模型)驱动的语音合成方案,支持语音克隆和语音创建功能,可通过命令行界面(CLI)和 Web UI 两种方式使用。用户可以根据需求调整语音的性别、音高、速度等参数,生成高质量的语音。该项目适用于多种场景,如有声读物制作、智能语音助手开发等。
字节跳动发布的AI编程神器IDE
Trae是一种自适应的集成开发环境(IDE),通过自动化和多元协作改变开发流程。利用Trae,团队能够更快速、精确地编写和部署代码,从而提高编程效率和项目交付速度。Trae具备上下文感知和代码自动完成功能,是提升开发效率的理想工具。
AI助力,做PPT更简单!
咔片是一款轻量化在线演示设计工具,借助 AI 技术,实现从内容生成到智能设计的一站式 PPT 制作服务。支持多种文档格式导入生成 PPT,提供海量模板、智能美化、素材替换等功能,适用于销售、教师、学生等各类人群,能高效制作出高品质 PPT,满足不同场景演示需求。
选题、配图、成文,一站式创作,让内容运营更高效
讯飞绘文,一个AI集成平台,支持写作、选题、配图、排版和发布。高效生成适用于各类媒体的定制内容,加速品牌传播,提升内容营销效果。
专业的AI公文写作平台,公文写作神器
AI 材料星,专业的 AI 公文写作辅助平台,为体制内工作人员提供高效的公文写作解决方案。拥有海量公文文库、9 大核心 AI 功能,支持 30 + 文稿类型生成,助力快速完成领导讲话、工作总结、述职报告等材料,提升办公效率,是体制打工人的得力写作神器。
OpenAI Agents SDK,助力开发者便捷使用 OpenAI 相关功能。
openai-agents-python 是 OpenAI 推出的一款强大 Python SDK,它为开发者提供了与 OpenAI 模型交互的高效工具,支持工具调用、结果处理、追踪等功能,涵盖多种应用场景,如研究助手、财务研究等,能显著提升开发效率,让开发者更轻松地利用 OpenAI 的技术优势。
高分辨率纹理 3D 资产生成
Hunyuan3D-2 是腾讯开发的用于 3D 资产生成的强大工具,支持从文本描述、单张图片或多视角图片生成 3D 模型,具备快速形状生成能力,可生成带纹理的高质量 3D 模型,适用于多个领域,为 3D 创作提供了高效解决方案。
一个具备存储、管理和客户端操作等多种功能的分布式文件系统相关项目。
3FS 是一个功能强大的分布式文件系统项目,涵盖了存储引擎、元数据管理、客户端工具等多个模块。它支持多种文件操作,如创建文件和目录、设置布局等,同时具备高效的事件循环、节点选择和协程池管理等特性。适用于需要大规模数据存储和管理的场景,能够提高系统的性能和可靠性,是分布式存储领域的优质解决方案。
最新AI工具、AI资讯
独家AI资源、AI项目落地
微信扫一扫关注公众号