/MindDiffuser

Primary LanguagePythonMIT LicenseMIT

MindDiffuser

This is the official code for the paper "MindDiffuser: Controlled Image Reconstruction from Human Brain Activity with Semantic and Structural Diffusion"(ACMMM2023)



Schematic diagram of MindDiffuser. 
(a) Decoders are trained to fit fMRI with averaged CLIP text embeddings 𝑐, CLIP image feature 𝑍𝑖𝐢𝐿𝐼𝑃, and VQ-VAE latent feature 𝑧.
(b) The two-stage image reconstruction process. In stage 1, an initial reconstructed image is generated using the decoded CLIP text feature 𝑐 and VQ-VAE latent feature 𝑧. In stage 2, the decoded CLIP image feature is used as a constraint to iteratively adjust 𝑐 and 𝑧 until the final reconstruction result matches the original image in terms of both semantic and structure.



A brief comparison of image reconstruction results.



Reconstruction results of MindDiffuser on multiple subjects



During the feature decoding process, we use L2-regularized linear regression model to automatically select voxels to fit three types
of feature: semantic feature 𝑐, detail feature 𝑧, and structural feature 𝑍𝐢𝐿𝐼𝑃. We ultilize pycortex to project the weights of each 
voxel in the fitted model onto the corresponding 3D coordinates in the visual cortex.

Preliminaries

This code was developed and tested with:

  • Python version 3.8.5
  • PyTorch version 1.11.0
  • A100 40G
  • The conda environment defined in environment_1.yml

Dataset

NSD dataset http://naturalscenesdataset.org/
Data preparation https://github.com/styvesg/nsd

Experiments

MindDiffuser

Model preparation

First, set up the conda enviroment as follows:

conda env create -f environment_1.yml  # create conda env
conda activate MindDiffuser          # activate conda env

Feature extraction 

cd your_folder
python Feature extractor/Semantic_feature_extraction.py
python Feature extractor/detail_extracttion.py
python Feature extractor/Structural_feature_extraction.py
python Feature extractor/Structural_feature_selection.py

Feature decoding 

cd your_folder
python Feature decoding/Semantic_feature_decoding.py
python Feature decoding/Structural_feature_decoding.py
python Feature decoding/detail_decoding.py

Image reconstruction

I will upload files such as features.npy,mask.npy, checkpoints of decoders, etc. to the checkpoints folder.

cd your_folder
python Image reconstruction/Reconstruction.py

Reproduce the results of "High-resolution image reconstruction with latent diffusion models from human brain activity"(CVPR2023)

After extracting and decoding the features, run the following code：

cd your_folder
python Reproduce Takagi's results/image_reconstruction.py

Reproduce the results of "Reconstruction of Perceived Images from fMRI Patterns and Semantic Brain Exploration using Instance-Conditioned GANs"

After configuring the environment and codes provided by Ozcelik, run the following codes:

cd your_folder
python Reproduce Ozcelik's results/extract_features.py
python Reproduce Ozcelik's results/train_regression.py
python Reproduce Ozcelik's results/reconstruct_images.py