See the following code:
# Model Use
This model is an class condition DDPM that can generate ultra sound images. It is based on the diffusion model from the [Towards Realistic Ultrasound Fetal Brain Imaging Synthesis](https://arxiv.org/abs/2304.03941) paper. The dataset used to train the model is this [FETAL_PLANES_DB dataset](https://zenodo.org/record/3904280). The classes that can be generated and the associated integer labels are
| Label | Class |
| -------- | ------- |
| 0 | Fetal abdomen |
| 1 | Fetal brain |
| 2 | Fetal femur |
| 3 | Fetal thorax |
| 4 | Maternal cervix |
| 5 | Other |
When generating images simply provide the label of you chosen class as an argument to the UNet.
Below you will see code that allows you to load this model, generate an image, and display it.
```python
# !pip install --upgrade diffusers transformers accelerate scipy ftfy safetensors
from diffusers import DDPMPipeline, DDIMPipeline, PNDMPipeline
import torch
import matplotlib.pyplot as plt
import numpy as np
# Are we using a GPU or CPU?
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# load model and scheduler
model_id = "harveymannering/xfetus-ddpm-v2"
ddpm = DDPMPipeline.from_pretrained(model_id)
ddpm.to(device)
# Generate a single image
x = torch.randn(1, 3, 128, 128).to(device) # noise
for i, t in enumerate(ddpm.scheduler.timesteps):
model_input = ddpm.scheduler.scale_model_input(x, t)
with torch.no_grad():
# Conditiong on the 'Fetal brain' class (with index 1)
class_label = torch.ones(1, dtype=torch.int64)
noise_pred = ddpm.unet(model_input, t, class_label.to(device))["sample"]
x = ddpm.scheduler.step(noise_pred, t, x).prev_sample
# Display image
plt.imshow(np.transpose(x[0].cpu().detach().numpy(), (1,2,0)) + 0.5)
```
# Example Outputs
The images from the original dataset as well as synthetic images. The following preprocessing/augmentation steps were applied to every image:
1. Random Horizontal Flip
2. Random Rotation (±45°)
3. Resize with Bicubic Interpolation
# Training Loss
Baseline was trained on images from the "Voluson E6" machine only. Training and validation loss are given below. Checkpoints were saved every 50 epochs and the best-performing checkpoint on the validation loss was at epoch 250. The model provided here is the checkpoint from epoch 250.
