条件潜在扩散模型 (Conditional Latent Diffusion Model) for MRI T1→T2 转换

基于PyTorch实现的条件潜在扩散模型，用于医学图像模态转换（T1 MRI → T2 MRI）。

项目特点

• 🏥 医学图像专用: 专门针对MRI T1到T2模态转换优化
• 🚀 高效架构: 使用VAE潜在空间降低计算复杂度
• 🎯 条件生成: 基于T1图像条件生成对应的T2图像
• 📊 完整评估: 包含PSNR、SSIM等医学图像质量评估指标
• ⚡ 快速推理: 支持DDIM采样器实现快速生成
• 🔧 易于使用: 提供完整的训练和推理pipeline

系统架构

T1 图像 → VAE编码器 → 潜在表示 → 条件扩散模型 → 生成潜在表示 → VAE解码器 → T2 图像

核心组件

1. VAE (变分自编码器)

   • 将256×256图像编码到32×32×4潜在空间
   • 大幅降低扩散模型的计算复杂度
   • 支持高质量图像重建

2. 条件UNet

   • 基于注意力机制的UNet架构
   • 交叉注意力层融合T1条件信息
   • 时间步嵌入支持扩散过程

3. 扩散模型

   • 支持DDPM和DDIM采样
   • 可配置的噪声调度策略
   • 稳定的训练过程

快速开始

一键运行演示

python quick_start.py --mode demo

完整流程（演示 + 训练 + 推理）

python quick_start.py --mode full --epochs 20

仅运行训练

python quick_start.py --mode train --epochs 50

仅运行推理

python quick_start.py --mode inference

安装依赖

pip install -r requirements.txt

或使用快速开始脚本自动安装：

python quick_start.py --mode demo

数据准备

1. 创建数据目录结构：

data/
├── T1/          # T1 MRI图像 (PNG格式)
│   ├── image001.png
│   ├── image002.png
│   └── ...
└── T2/          # 对应的T2 MRI图像 (PNG格式)
    ├── image001.png
    ├── image002.png
    └── ...

2. 确保T1和T2图像文件名一一对应
3. 图像格式：PNG，灰度图，推荐分辨率256×256

训练

1. 完整训练流程

python train.py --mode full --epochs 100 --lr 1e-4

2. 分阶段训练

第一阶段：训练VAE

python train.py --mode vae --epochs 50 --lr 1e-4 --kl_weight 1e-6

第二阶段：训练扩散模型

python train.py --mode diffusion --epochs 100 --lr 1e-4 --resume_vae checkpoints/vae_best.pth

训练参数说明

• --mode: 训练模式 (vae/diffusion/full)
• --epochs: 训练轮数
• --lr: 学习率
• --batch_size: 批次大小
• --kl_weight: VAE的KL散度损失权重
• --data_split: 训练/验证集分割比例
• --resume_vae: 恢复VAE检查点
• --resume_diffusion: 恢复扩散模型检查点

推理

1. 单图像推理

python inference.py \
    --mode single \
    --vae_path checkpoints/vae_best.pth \
    --diffusion_path checkpoints/diffusion_best.pth \
    --input path/to/t1_image.png \
    --output path/to/generated_t2.png

2. 批量推理

python inference.py \
    --mode batch \
    --vae_path checkpoints/vae_best.pth \
    --diffusion_path checkpoints/diffusion_best.pth \
    --input data/T1/ \
    --output outputs/generated/

3. 模型评估

python inference.py \
    --mode evaluate \
    --vae_path checkpoints/vae_best.pth \
    --diffusion_path checkpoints/diffusion_best.pth \
    --input data/T1/ \
    --t2_dir data/T2/ \
    --output outputs/evaluation/

推理参数说明

• --use_ddim: 使用DDIM采样器（默认开启）
• --num_inference_steps: 推理步数（默认50步）
• --num_samples: 评估时的样本数量限制

配置文件

主要配置在 config.py 中：

class Config:
    # 数据配置
    IMAGE_SIZE = 256
    BATCH_SIZE = 8
    
    # 模型配置
    LATENT_DIM = 4
    TIMESTEPS = 1000
    
    # 训练配置
    LEARNING_RATE = 1e-4
    NUM_EPOCHS = 100

评估指标

系统提供以下评估指标：

• PSNR (峰值信噪比): 衡量图像重建质量
• SSIM (结构相似性): 评估图像结构保持程度

评估结果会保存为JSON格式：

{
  "psnr_mean": 28.45,
  "psnr_std": 2.31,
  "ssim_mean": 0.892,
  "ssim_std": 0.045,
  "num_samples": 100
}

输出文件

训练过程中会生成：

• checkpoints/: 模型检查点
• logs/: 训练日志和TensorBoard文件
• outputs/samples/: 训练过程中的生成样本
• outputs/reconstructions/: VAE重建结果
• outputs/comparisons/: 对比图像

模型架构细节

VAE架构

• 编码器: 4层下采样，ResNet块 + 注意力机制
• 解码器: 4层上采样，对称结构
• 潜在维度: 4通道，32×32分辨率

UNet架构

• 输入/输出: 4通道潜在表示
• 时间嵌入: 正弦位置编码
• 条件融合: 交叉注意力机制
• 注意力层: 在32×32、16×16、8×8分辨率

扩散过程

• 前向过程: 线性或余弦噪声调度
• 反向过程: 学习噪声预测
• 采样: 支持DDPM和DDIM

性能优化建议

1. 内存优化:

   • 使用梯度检查点
   • 调整批次大小
   • 启用混合精度训练

2. 训练加速:

   • 使用多GPU训练
   • 预训练VAE权重
   • 调整学习率调度

3. 质量提升:

   • 增加训练数据
   • 调整损失函数权重
   • 使用数据增强

故障排除

常见问题

1. CUDA内存不足:

   # 减小批次大小
   python train.py --batch_size 4

2. 数据加载错误:

   • 检查图像文件格式和命名
   • 确保T1/T2图像配对正确

3. 训练不稳定:

   • 降低学习率
   • 调整KL散度权重
   • 使用梯度裁剪

扩展功能

• 支持3D MRI数据
• 多模态条件生成
• 实时推理优化
• Web界面部署
• 医学图像预处理工具

引用

如果您使用了本项目，请考虑引用相关论文：

@article{rombach2022high,
  title={High-resolution image synthesis with latent diffusion models},
  author={Rombach, Robin and Blattmann, Andreas and Lorenz, Dominik and Esser, Patrick and Ommer, Bj{\"o}rn},
  journal={CVPR},
  year={2022}
}

许可证

MIT License

联系方式

如有问题或建议，请提交Issue或Pull Request。