# Code for Neuro-Vision to Language: Image Reconstruction and Language-Enabled Interaction via Brain Recordings

## Environment

In addition to the packages listed in `requirements.txt`, the source code from [LLaVA](https://github.com/haotian-liu/LLaVA) is required for multimodal instruction fine-tuning.
Also, [StableDiffusionReconstruction](https://github.com/yu-takagi/StableDiffusionReconstruction) is necessary to construct the NSD dataset.

## Prepare Datasets

First, download the NSD dataset from [here](https://cvnlab.slite.page/p/M3ZvPmfgU3/General-Information) and place it appropriately.
Then, run `prepare_datasets/make_subjmri.py` to generate fMRI data for different subjects. This step requires the `StableDiffusionReconstruction` environment.
Further documentation assumes that the processed data is located at `/mnt/NSD_datasets/datasets/nsd`.

### Structure of the dataset

The final folder structure is as follows, where `datasets` is the root directory and `nsd` refers to the `Natural Scenes Dataset`.

It contains:
1. `fmris` - fMRI data, with subdirectories for different `subject`.
2. `images` - Image data, i.e., the images each subject viewed.
3. `nsd_captions.json` - Image descriptions generated using BLIP2.
4. `nsd_coco_captions.json` - Descriptions from the COCO dataset, if available.
5. `nsd_instances.json` - Instance details for each image in COCO format, but with coordinates for bounding box vertices.
6. `nsd_gpt_conversation` - Dialogue data generated from `nsd_captions.json`, `nsd_coco_captions.json`, and `nsd_instances.json`.
7. `sft_data_tr.json` and `sft_data_te.json` - Data for training and testing, formatted for direct use in training.
8. `vision_embeds` - Feature vectors for the images, used for aligning with the vision features during the training of the fMRI data.

### Prepare the dataset

First, prepare the `fMRI` and `Image` data, which correspond to the images subjects viewed and their brain activity.

#### fMRI

The first layer under the `fMRI` path is different `subj` data, i.e., different subjects should be separated.

##### Area
If the dataset provides segmentation for different brain areas, generate the `area` path.
The structure is as follows:

```
area
|-- nsd_early_betas_mean.npy
|-- nsd_early_betas_std.npy
|-- nsd_early_betas_te.npy
|-- nsd_early_betas_tr.npy
...
```


All test data for the same brain area are saved in one `npy` file, named `{dataset}_{area}_betas_{tr or te}.npy`,
where `tr` denotes training set, and `te` denotes test set.
The array shape is `(tr/te_samples, area_dims)`, meaning the betas values for the brain area are flattened into a one-dimensional vector.

Also include mean and standard deviation data for normalization during sequential reading, named `nsd_early_betas_mean.npy` and `nsd_early_betas_std.npy`.
The array shape for these is `(area_dims,)`.

##### Whole

**Must** include whole-brain fMRI data. The structure is as follows:


```json
{
    "train": [
      61882,
      828,
      67573,
      16020,
      40422,
      51517,
      62325,
      50610,
      ...
    ],
    "val": [
      61882,
      828,
      67573,
      16020,
      40422,
      51517,
      62325,
      50610,
      ...
    ]
}
```

For instance, `fmri2image['train'][0] = 61882`, meaning the first sample in the training set corresponds to image index `61882`.

Then we complete the preparation of the `fMRI` data.

#### Image
The indexing for `Image` does not need to match that of `fMRI`, but correspondence can be established through the `nsd_fmri2image.json` file. This file's structure is simple, consisting of a directory with paths to images.

```
images
...
|-- nsd_image_072981.png
|-- nsd_image_072982.png
|-- nsd_image_072983.png
|-- nsd_image_072984.png
|-- nsd_image_072985.png
|-- nsd_image_072986.png
|-- nsd_image_072987.png
|-- nsd_image_072988.png
|-- nsd_image_072989.png
|-- nsd_image_072990.png
|-- nsd_image_072991.png
|-- nsd_image_072992.png
|-- nsd_image_072993.png
...
```


Images are named in the format `{dataset}_image_{index:06}.png`, where `index` is a six-digit number representing the image index.

When saving images using `PIL`, the parameter `lossless=True` should be passed to ensure the quality of the images and avoid [electronic patina](https://zhuanlan.zhihu.com/p/590753356).

### Generate the dataset

With the aforementioned parts in place, you can start generating the complete dataset. Each generation script can directly accept the dataset name, meaning you can run `python script.py --dataset {dataset}` to generate the dataset. If the preliminary work is done correctly, it should work seamlessly.

#### Generate the captions

This mainly uses the blip2 model to generate descriptions for the images. The output is saved in `nsd_captions.json`.

Use `generate_captions.py` to produce these.

#### Generate the instances

This mainly uses the `DETR` model to generate categories and bounding boxes for instances in the images. The output is saved in `nsd_instances.json`.

Use `objection_detection.py` to generate these.

#### Generate the vision embeddings

This involves the `CLIP` and `AutoencoderKL` models to generate feature vectors for the images. The output folders are `vision_embeds` and `vae_embeds`.

Use `generate_clip_embeds.py` to produce these.

#### Generate the gpt conversation

With the `instances` and `captions` ready, `GPT` can generate conversations. The output folder is `nsd_gpt_conversation`.

Use `generate_conversation.py` for this. It needs to be run four times to generate four types of dialogues: multi-turn conversations, detailed descriptions, and complex reasoning.

The parameters to pass are `--prompt_type {complex_reasoning, detail_description, conversation}`.

#### Generate the sft data

With all the above data ready, we can generate a dataset for multimodal SFT compliant with the [format](https://github.com/haotian-liu/LLaVA/blob/main/docs/Finetune_Custom_Data.md) described by `llava`. This includes `sft_data_tr.json` and `sft_data_te.json`, which will also be provided with the supplementary materials.

## Train fMRI Extractor

Once the dataset is prepared, you can begin training the `fMRI` `vision tower`. First, you'll need the [LLaVA](https://github.com/haotian-liu/LLaVA) source code, then supplement the downloaded source with files from the supplementary material for fMRI compatibility.

Then, run `pretrain_vit3d` to pretrain the fMRI vision tower.
```bash
    python pretrain_vit3d.py --num-layers 16 --patch-size 14 --subject all --select-region nsdgeneral --vae
```

The json file used to train the extractor is provided in `data/pretrain.json`.

## Instruction Fine-tuning

After that, you can proceed with two phases of instruction fine-tuning. 
The first phase is fine-tuning the fMRI vision tower, and the second phase is joint fine-tuning of the vision tower and LLMs. 
The scripts used for fine-tuning can be found in `scripts/brain_llava/pretrain_adapter.sh` and `scripts/brain_llava/llava_adapter.sh`. 
Note that you need to replace `--vision_tower` with the path to your pretrained vision tower. Due to the size limits of supplementary materials, we will provide the fine-tuned model after the manuscript review.

The json files used for instruction fine-tuning are provided in `data/sft_all_tr.json` and `data/sft_all_te.json`.

## Visual Reconstruction

Using the brief descriptions of visual stimuli and the vision tower's predicted visual embeddings, visual reconstruction can be performed.
We provide a series of scripts for this step, which can be found in `./tools/evaluate`.

1. First run `fmri2embeds.py` to generate the output of the fMRI`vision tower`.
2. Run `generate_captions_from_llava.py` to generate brief descriptions, detailed descriptions, complex reasoning and other results.
3. Then run `embeds2images.py` to generate the reconstructed image.


## Evaluation

For visual reconstruction, we use the same evaluation script as [MindEye](https://github.com/MedARC-AI/fMRI-reconstruction-NSD).

For tasks such as brain captions, we use `pycocoevalcap` for evaluation.