File size: 1,738 Bytes
ca205bf cd0724e ca205bf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
---
license: mit
tags:
- pytorch
- diffusers
- unconditional-image-generation
- diffusion-models-class
- medical-imaging
- brain-mri
- multiple-sclerosis
---
# Brain MRI Synthesis with DDPM (64x64)
This model is a diffusion-based model for unconditional image generation of **brain MRI FLAIR slices** of size **64x64 pixels**.
The model was trained using the [DDPM](https://arxiv.org/abs/2006.11239) architecture, with attention mechanisms in the middle of the U-Net.
It is trained from scratch on a dataset of brain MRI slices, specifically designed for generating synthetic brain images.
## Training Details
- **Architecture:** DDPM (Denoising Diffusion Probabilistic Model)
- **Resolution:** 64x64 pixels
- **Dataset:** Lesion2D VH splitted (FLAIR MRI slices) (70% of the dataset)
- **Channels:** 1 (grayscale, FLAIR modality)
- **Epochs:** 50
- **Batch size:** 32
- **Optimizer:** AdamW with learning rate of `1.0e-4`
- **Scheduler:** Cosine with 500 warm-up steps
- **Gradient Accumulation:** 1 steps
- **Mixed Precision:** No
- **Hardware:** Trained on **one NVIDIA GeForce GTX 1080 Ti GPU of 12GB**
- **Memory Consumption:** Around **7 GB** during training
## U-Net Architecture
- **Down Blocks:** [DownBlock2D, DownBlock2D, AttnDownBlock2D, DownBlock2D]
- **Up Blocks:** [UpBlock2D, AttnUpBlock2D, UpBlock2D]
- **Layers per Block:** 2
- **Block Channels:** [128, 128, 256, 512]
## Usage
You can use the model directly with the `diffusers` library:
```python
from diffusers import DDPMPipeline
import torch
# Load the model
pipeline = DDPMPipeline.from_pretrained("benetraco/brain_ddpm_64")
pipeline.to("cuda") # or "cpu"
# Generate an image
image = pipeline(batch_size=1).images[0]
# Display the image
image.show()
|