--- a/README.md
+++ b/README.md
@@ -1,231 +1,218 @@
-<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}
-}
-```
+<div align="center">
+
+<h2 style="border-bottom: 1px solid lightgray;">Visual Decoding and Reconstruction via EEG Embeddings with Guided Diffusion</h2>
+
+
+<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}
+}
+```