[ICASSP 2025] Guess What I Think: Streamlined EEG-to-Image Generation with Latent Diffusion Models

Project page here 🚀

If you like our project, please give us a star ⭐ on GitHub for latest update.

Eleonora Lopez, Luigi Sigillo, Federica Colonnese, Massimo Panella and Danilo Comminiello

ISPAMM Lab and NESYA Lab , Sapienza University of Rome

📰 News

• [2024.12.25] 🔥🔥🔥 Request the checkpoints available here!
• [2024.12.25] Code is available now! Welcome to watch 👀 this repository for the latest updates.
• [2024.12.20] The paper has been accepted for presentation at ICASSP 2025 🎉!
• [2024.09.17] The paper has been published on Arxiv 🎉. The pdf version is available here!

😮 Highlights

💡 Revolutionizing Brain-Computer Interfaces (BCIs)

Our work bridges the gap between brain signals and visual understanding by tackling the challenge of generating images directly from EEG signals—a low-cost, non-invasive, and portable neuroimaging modality suitable for real-time applications.

🔥 State-of-the-Art EEG-to-Image Generation

We propose a streamlined framework leveraging the ControlNet adapter to condition a latent diffusion model (LDM) through EEG signals. Extensive experiments and ablation studies on popular benchmarks show our method outperforms state-of-the-art models while requiring only minimal preprocessing and fewer components.

👀 A Minimal yet Powerful Approach

Unlike existing methods that demand heavy preprocessing, complex architectures, and additional components like captioning models, our approach is efficient and straightforward. This enables a new frontier in real-time BCIs, advancing tasks like visual cue decoding and future neuroimaging applications.

🚀 Main Results

For more evaluation, please refer to our paper for details.

How to run experiments 💻

Building Environment

conda create --name=gwit python=3.9
conda activate gwit

pip install src/diffusers

pip install transformers accelerate xformers==0.0.16 wandb numpy==1.26.4 datasets torchvision==0.14.1 scikit-learn

Train

To launch the training of the model, you can use the following command, you need to change the output_dir:

accelerate launch src/gwit/train_controlnet.py --caption_from_classifier --subject_num=4 --pretrained_model_name_or_path=stabilityai/stable-diffusion-2-1-base --output_dir=output/model_out_CVPR_SINGLE_SUB_CLASSIFIER_CAPTION --dataset_name=luigi-s/EEG_Image_CVPR_ALL_subj --conditioning_image_column=conditioning_image --image_column=image --caption_column=caption --resolution=512 --learning_rate=1e-5 --train_batch_size=8 --num_train_epochs=50 --tracker_project_name=controlnet --enable_xformers_memory_efficient_attention --checkpointing_steps=1000 --validation_steps=500 --report_to wandb --validation_image ./using_VAL_DATASET_PLACEHOLDER.jpeg --validation_prompt "we are using val dataset hopefuly"

You can change the dataset using one of, with the dataset_name parameter:

• luigi-s/EEG_Image_CVPR_ALL_subj
• luigi-s/EEG_Image_TVIZ_ALL_subj

Generate

Request access to the pretrained models from Google Drive.

To launch the generation of the images from the model, you can use the following command, you need to change the output_dir:

python src/gwit/validate_controlnet.py --controlnet_path=output/model_out_CVPR_SINGLE_SUB_CLASSIFIER_CAPTION/checkpoint-24000/controlnet/ --caption --single_image_for_eval --guess

Evaluation

Request access to the pretrained models from Google Drive.

To launch the testing of the model, you can use the following command, you need to change the output_dir:

python gwit/evaluation/evaluate.py --controlnet_path=output/model_out_CVPR_SINGLE_SUB_CLASSIFIER_CAPTION/checkpoint-24000/controlnet/ --caption --single_image_for_eval --guess

Dataset

The dataset used are hosted on huggingface:

• ImageNetEEG
• Thoughtviz

Cite

Please cite our work if you found it useful:

@misc{lopez2024guessithinkstreamlined,
      title={Guess What I Think: Streamlined EEG-to-Image Generation with Latent Diffusion Models}, 
      author={Eleonora Lopez and Luigi Sigillo and Federica Colonnese and Massimo Panella and Danilo Comminiello},
      year={2024},
      eprint={2410.02780},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2410.02780}, 
}

Star History

Acknowledgement

This project is based on diffusers. Thanks for their awesome work.