# MindEye v2 图像重建项目

该项目基于功能性磁共振成像(fMRI)数据实现从大脑活动到图像的重建。通过深度学习模型分析人脑在观看图像时的神经活动模式，并尝试重建出被试者看到的图像内容。

## 项目结构

```
.
├── mindeye2_src/                 # 核心模型代码库
│   ├── generative_models/        # 生成模型相关代码
│   ├── sgm/                      # 稳定扩散模型组件
│   └── 其他模型和工具代码
├── dataset/                      # 数据集目录（需用户准备）
│   ├── betas_all_subj0{}_fp32_renorm.hdf5  # fMRI数据
│   ├── clip_embeddings.hdf5      # CLIP嵌入特征
│   └── wds/                      # WebDataset格式的训练数据
├── train_vae_eye.py              # VAE模型训练脚本
├── train_prior_vae_eye.py        # Prior模型训练脚本
├── mindeye_recon.py              # 图像重建脚本
├── models.py                     # 自定义模型定义
└── 其他辅助脚本
```

## 功能模块

### 1. VAE模型训练 (`train_vae_eye.py`)
- 使用变分自编码器(VAE)对fMRI数据进行编码和解码
- 提取fMRI数据的关键特征表示
- 保存训练好的VAE模型供后续使用

### 2. Prior模型训练 (`train_prior_vae_eye.py`)
- 训练先验网络，建立CLIP图像嵌入与fMRI特征之间的映射关系
- 使用扩散模型优化CLIP特征与fMRI特征的对齐

### 3. 图像重建 (`mindeye_recon.py`)
- 利用训练好的VAE和Prior模型从fMRI数据重建图像
- 支持生成不同风格的重建图像

## 使用方法

### 环境准备
```bash
pip install -r requirements.txt
```

### 数据准备
在项目根目录下创建`dataset`目录，并放置以下数据文件：
- `betas_all_subj0{}_fp32_renorm.hdf5`: fMRI神经活动数据
- `clip_embeddings.hdf5`: 对应图像的CLIP嵌入特征
- `wds/`: WebDataset格式的训练数据

### 模型训练
1. 训练VAE模型:
```bash
python train_vae_eye.py
```

2. 训练Prior模型:
```bash
python train_prior_vae_eye.py
```

### 图像重建
```bash
python mindeye_recon.py
```

## 主要依赖
- PyTorch
- H5Py (处理科学数据)
- WebDataset (大规模数据集处理)
- TensorBoard (训练可视化)
- CLIP (图像文本表征)
- Stable Diffusion (图像生成)

## 注意事项
- 所有脚本默认使用CUDA设备进行训练，确保有可用的GPU资源
- 训练参数可根据具体需求调整
- 数据路径和被试者编号可在脚本中配置