--- a
+++ b/README.md
@@ -0,0 +1,231 @@
+<div align="center">
+
+<h2 style="border-bottom: 1px solid lightgray;">Visual Decoding and Reconstruction via EEG Embeddings with Guided Diffusion</h2>
+
+<!-- Badges and Links Section -->
+<div style="display: flex; align-items: center; justify-content: center;">
+
+<p align="center">
+  <a href="#">
+  <p align="center">
+    <a href='https://arxiv.org/pdf/2403.07721'><img src='http://img.shields.io/badge/Paper-arxiv.2403.07721-B31B1B.svg'></a>
+    <a href='https://huggingface.co/datasets/LidongYang/EEG_Image_decode/tree/main'><img src='https://img.shields.io/badge/EEG Image decode-%F0%9F%A4%97%20Hugging%20Face-blue'></a>
+  </p>
+</p>
+
+
+</div>
+
+<br/>
+
+</div>
+
+<!-- 
+<img src="bs=16_test_acc.png" alt="Framework" style="max-width: 90%; height: auto;"/> -->
+<!-- 
+<img src="test_acc.png" alt="Framework" style="max-width: 90%; height: auto;"/> -->
+
+<!-- As the training epochs increases, the test set accuracy of different methods. (Top: batchsize is 16. Bottom: batchsize is 1024) -->
+
+<!-- 
+<img src="temporal_analysis.png" alt="Framework" style="max-width: 90%; height: auto;"/>
+Examples of growing window image reconstruction with 5 different random seeds. -->
+
+
+<img src="fig-framework.png" alt="Framework" style="max-width: 100%; height: auto;"/>
+Framework of our proposed method.
+
+
+
+
+<!--  -->
+<img src="fig-genexample.png" alt="fig-genexample" style="max-width: 90%; height: auto;"/>  
+
+Some examples of using EEG to reconstruct stimulus images.
+
+
+## News:
+- [2024/09/26] Our paper is accepted to **NeurIPS 2024**.
+- [2024/09/25] We have updated the [arxiv](https://arxiv.org/abs/2403.07721) paper.
+- [2024/08/01] Update scripts for training and inference in different tasks.
+- [2024/05/19] Update the dataset loading scripts.
+- [2024/03/12] The [arxiv](https://arxiv.org/abs/2403.07721) paper is available.
+
+
+<!-- ## Environment setup -->
+<h2 style="border-bottom: 1px solid lightgray; margin-bottom: 5px;">Environment setup</h2>
+
+Run ``setup.sh`` to quickly create a conda environment that contains the packages necessary to run our scripts; activate the environment with conda activate BCI.
+
+
+
+
+```
+. setup.sh
+```
+You can also create a new conda environment and install the required dependencies by running
+```
+conda env create -f environment.yml
+conda activate BCI
+
+pip install wandb
+pip install einops
+```
+Additional environments needed to run all the code:
+```
+pip install open_clip_torch
+
+pip install transformers==4.28.0.dev0
+pip install diffusers==0.24.0
+
+#Below are the braindecode installation commands for the most common use cases.
+pip install braindecode==0.8.1
+```
+<!-- ## Quick training and test  -->
+<h2 style="border-bottom: 1px solid lightgray; margin-bottom: 5px;">Quick training and test</h2>
+
+If you want to quickly reproduce the results in the paper, please download the relevant ``preprocessed data`` and ``model weights`` from [Hugging Face](https://huggingface.co/datasets/LidongYang/EEG_Image_decode) first.
+#### 1.Image Retrieval
+We provide the script to learn the training strategy of EEG Encoder and verify it during training. Please modify your data set path and run:
+```
+cd Retrieval/
+python ATMS_retrieval.py --logger True --gpu cuda:0  --output_dir ./outputs/contrast
+```
+We also provide the script for ``joint subject training``, which aims to train all subjects jointly and test on a specific subject:
+```
+cd Retrieval/
+python ATMS_retrieval_joint_train.py --joint_train --sub sub-01 True --logger True --gpu cuda:0  --output_dir ./outputs/contrast
+```
+
+Additionally, replicating the results of other methods (e.g. EEGNetV4) by run
+```
+cd Retrieval/
+contrast_retrieval.py --encoder_type EEGNetv4_Encoder --epochs 30 --batch_size 1024
+```
+
+#### 2.Image Reconstruction
+We provide quick training and inference scripts for ``clip pipeline`` of visual reconstruction. Please modify your data set path and run zero-shot on 200 classes test dataset:
+```
+# Train and generate eeg features in Subject 8
+cd Generation/
+python ATMS_reconstruction.py --insubject True --subjects sub-08 --logger True \
+--gpu cuda:0  --output_dir ./outputs/contrast
+```
+
+```
+# Reconstruct images in Subject 8
+Generation_metrics_sub8.ipynb
+```
+
+We also provide scripts for image reconstruction combined ``with the low level pipeline``.
+```
+cd Generation/
+
+# step 1: train vae encoder and then generate low level images
+train_vae_latent_512_low_level_no_average.py
+
+# step 2: load low level images and then reconstruct them
+1x1024_reconstruct_sdxl.ipynb
+```
+
+
+We provide scripts for caption generation combined ``with the semantic level pipeline``.
+```
+cd Generation/
+
+# step 1: train feature adapter
+image_adapter.ipynb
+
+# step 2: get caption from eeg latent
+GIT_caption_batch.ipynb
+
+# step 3: load text prompt and then reconstruct images
+1x1024_reconstruct_sdxl.ipynb
+```
+
+To evaluate the quality of the reconstructed images, modify the paths of the reconstructed images and the original stimulus images in the notebook and run:
+```
+#compute metrics, cited from MindEye
+Reconstruction_Metrics_ATM.ipynb
+```
+
+<!-- ## Data availability -->
+<h2 style="border-bottom: 1px solid lightgray; margin-bottom: 5px;">Data availability</h2>
+
+We provide you with the ``preprocessed EEG`` and ``preprocessed MEG`` data used in our paper at [Hugging Face](https://huggingface.co/datasets/LidongYang/EEG_Image_decode), as well as the raw image data.
+
+
+Note that the experimental paradigms of the THINGS-EEG and THINGS-MEG datasets themselves are different, so we will provide images and data for the two datasets separately.
+
+You can also download the relevant THINGS-EEG data set and THINGS-MEG data set at osf.io.
+
+The raw and preprocessed EEG dataset, the training and test images are available on [osf](https://osf.io/3jk45/).
+- ``Raw EEG data:`` `../project_directory/eeg_dataset/raw_data/`.
+- ``Preprocessed EEG data:`` `../project_directory/eeg_dataset/preprocessed_data/`.
+- ``Training and test images:`` `../project_directory/image_set/`.
+
+
+The raw and preprocessed MEG dataset, the training and test images are available on [OpenNEURO](https://openneuro.org/datasets/ds004212/versions/2.0.0).
+
+
+
+
+
+<!-- ## EEG/MEG preprocessing -->
+<h2 style="border-bottom: 1px solid lightgray; margin-bottom: 5px;">EEG/MEG preprocessing</h2>
+
+
+Modify your path and execute the following code to perform the same preprocessing on the raw data as in our experiment:
+```
+cd EEG-preprocessing/
+python EEG-preprocessing/preprocessing.py
+```
+
+```
+cd MEG-preprocessing/
+MEG-preprocessing/pre_possess.ipynb
+```
+Also You can get the data set used in this project through the BaiduNetDisk [link](https://pan.baidu.com/s/1-1hgpoi4nereLVqE4ylE_g?pwd=nid5) to run the code.
+
+## TODO
+- [√] Release retrieval and reconstruction scripts.
+- [√] Update training scripts of reconstruction pipeline.
+- [ ] Adding validation sets improves performance evaluation accuracy.
+
+
+
+<!-- ## Acknowledge -->
+<h2 style="border-bottom: 1px solid lightgray; margin-bottom: 5px;">Acknowledge</h2>
+
+1.Thanks to Y Song et al. for their contribution in data set preprocessing and neural network structure, we refer to their work:</br>"[Decoding Natural Images from EEG for Object Recognition](https://arxiv.org/pdf/2308.13234.pdf)".</br> Yonghao Song, Bingchuan Liu, Xiang Li, Nanlin Shi, Yijun Wang, and Xiaorong Gao. 
+
+2.We also thank the authors of [SDRecon](https://github.com/yu-takagi/StableDiffusionReconstruction) for providing the codes and the results. Some parts of the training script are based on [MindEye](https://medarc-ai.github.io/mindeye/) and [MindEye2](https://github.com/MedARC-AI/MindEyeV2). Thanks for the awesome research works.
+
+3.Here we provide our THING-EEG dataset cited in the paper:</br>"[A large and rich EEG dataset for modeling human visual object recognition](https://www.sciencedirect.com/science/article/pii/S1053811922008758?via%3Dihub)".</br>
+Alessandro T. Gifford, Kshitij Dwivedi, Gemma Roig, Radoslaw M. Cichy.
+
+
+4.Another used THINGS-MEG data set provides a reference:</br>"[THINGS-data, a multimodal collection of large-scale datasets for investigating object representations in human brain and behavior.](https://elifesciences.org/articles/82580.pdf)".</br> Hebart, Martin N., Oliver Contier, Lina Teichmann, Adam H. Rockter, Charles Y. Zheng, Alexis Kidder, Anna Corriveau, Maryam Vaziri-Pashkam, and Chris I. Baker.
+
+
+
+<!-- ## Citation -->
+<h2 style="border-bottom: 1px solid lightgray; margin-bottom: 5px;">Citation</h2>
+
+```bibtex
+@inproceedings{
+li2024visual,
+title={Visual Decoding and Reconstruction via {EEG} Embeddings with Guided Diffusion},
+author={Dongyang Li and Chen Wei and Shiying Li and Jiachen Zou and Quanying Liu},
+booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
+year={2024},
+url={https://openreview.net/forum?id=RxkcroC8qP}
+}
+
+@article{li2024visual,
+  title={Visual Decoding and Reconstruction via EEG Embeddings with Guided Diffusion},
+  author={Li, Dongyang and Wei, Chen and Li, Shiying and Zou, Jiachen and Liu, Quanying},
+  journal={arXiv preprint arXiv:2403.07721},
+  year={2024}
+}
+```