# whole_brain_encoder

# Run inference on your images using the ensemble model

The ensemble model is a combination of many individual models, each trained on a particular subject, encoder layer, and hemisphere. It uses the validation split to generate a model confidence value for each voxel, so it should be evaluated on the test split or unseen data.

The ensemble model used in the submission includes weights for two runs each for
- `lh` and `rh` hemispheres
- encoder layers 1, 3, 5, 7 from the dino backbone.

The model weights are not provided in this submission in order to maintain anonymity. Follow the training instructions below in [Reproducing the checkpoints](#reproducing-the-checkpoints) to train the checkpoints yourself. The expected directory structure is: `checkpoints/nsd_test/dinov2_q_transformer/subj_{subj_num:02}/enc_{enc_layer}/run_{run_num}/{hemi}`.

## Environment setup

### Conda environment for model training and inference

Set up the conda environments:

```bash
conda env create -f env/{env}.yml
```

# Reproducing the checkpoints

## Required setup

In `utils/args.py`, modify the following paths in the default arguments:

1. `--data_dir`: Directory containing metadata and neural data files. These are from NSD.
1. `--imgs_dir`: Directory containing NSD image data. These are from NSD.

## Generating the brain parcels

```bash
conda activate xformers
cd img_gen
python roi_selection --subj ${subj} --hemi ${hemi}
```

This will select parcels for further experimentation based on the criteria described in the paper. The data for selected parcels is stored in `img_gen/schaefer/candidate_parcels/subj_{subj_num:02}/{hemi}/{parcel_dir}.npy`.

## Model training

Set up the [conda environment for model training and inference](#conda-environment-for-model-training-and-inference).

To train a single model, run:

```bash
conda activate xformers
python main.py --subj $SUBJECT --enc_output_layer $layer --run $RUN_ID --hemi $HEMI --backbone_arch $BACKBONE_ARCH
```

Backbone arch options include:
- radio_v2.5-h
- dinov2_q
- dinov2_q_large
- clip

# Running inference

Follow the example in `tutorials/test_wrapper.ipynb`.

To run inference on a folder (e.g., ImageNet), run:

```bash
conda activate xformers
python brain_encoder_wrapper.py --subj ${subject} --split folder --target_dir ${target_dir} --save_path ${save_path}  --save ${save} --backbone_arch ${backbone_arch}
```

To save attention masp: add flag --save_attention True

save options:
- all: save all full brain predictions
- parcel: save predictions only for parcels selected for further experimentation
- nsd_labeled: save predictions only for the labeled area
- none: do not save any predictions

# Backbone model comparison

Train ensemble models as described in [Run inference on your images using the ensemble model](#run-inference-on-your-images-using-the-ensemble-model) for backbones CLIP, dinov2_q, dinov2_q_large, and radio_v2.5-h.

```bash
python brain_encoder_wrapper.py --subj ${subject} --split test
```

Pearson's correlation results will be saved in `results/schaefer/enc_1_3_5_7_run_1_2/subj_{subj_num:02}/val_corr_avg.txt`.

# Superstimulus generation

## Top NSD images

```python
from fetch import top_NSD_imgs
nsd_test_imgs, top_nsd_idxs, top_nsd_model_idxs = top_NSD_imgs(subj, hemi, split="test")
```
nsd_test_imgs: list of NSDImage objects which include attributes
- nsd_test_imgs.image: image as an array
- nsd_test_imgs.activation[parcel_dir]: ground truth activation for parcel_dir
- nsd_test_imgs.model_activations[parcel_dir]: model activations for parcel_dir

## BrainDIVE generated superstimuli

```bash
conda activate braindivexf
python generate_imgs.py --subj ${subj} --hemi ${hemi} --parcel_dir ${parcel_dir} --num_imgs_to_generate ${num_imgs}
```

Images should be generated for each parcel_dir in `img_gen/schaefer/candidate_parcels/subj_{subj_num:02}/{hemi}/{parcel_dir}.npy`. The generated images are saved in `./images`.

Images can be rerated using the inference pipeline to further select maximally parcel-activating images.

```python
from fetch import top_generated_imgs
top_generated_imgs, top_img_paths, top_img_activations = top_generated_imgs(subj, hemi, parcel_dir, max_num_imgs=32)
```
top_generated_imgs: PIL Image
top_img_paths: list of paths to the generated images
top_img_activations: list of model activations for the generated images

## ImageNet encoder-selected superstimuli

Use the inference pipeline on all of ImageNet training set. Then choose top images to further analyze. Results should be stored in `results/schaefer/enc_1_3_5_7_run_1_2/imagenet/subj_01`.

# Evaluating our labels

```python
conda activate clip
cd clip_hypothesis
python embed_clip_patch.py # generates clip embeddings for all NSD images and top ImageNet/generated images
python cross_subj_rankcorr_stats_test.py
```

Saves `clip_hypothesis/cross_subj_p/clip_retrieval_rankcorr_{img_type}_{split}.npy`.

## Creating table 2 and 3

This only applies to the test split, since the training split was used to train the model.

For a given parcel_dir, read the appropriate .npy file as data and comapre data["intermediates"][parcel_dir]["retrieved_corr"] with null distribution data["intermediates"][parcel_dir]["null_corr"]. Also fetch the top 32 NSD train images using `fetch.top_imgnet_imgs(subj, hemi, parcel_dir)`, then embed these images into CLIP space and average. Follow the example in clip_hypothesis/cross_subj_rankcorr_stats_test.py to fetch the test images and compute the rank correlation.

## Creating tables 4 and 5

This applies to the training split, since in cross-subject analyses, data from all subjects other than the target are used to predict the data for the target. The images presented in the training split are unique for all subjects.

Follow the same procedure as above. For the top 32 NSD train images, average across the subjects using to predict the target subject.