Update README.md

README.md

@@ -1,36 +1,62 @@
 See the following code:
 
-## Model Use
+# 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 diffusers
+# !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)  # you can replace DDPMPipeline with DDIMPipeline or PNDMPipeline for faster inference
+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 tqdm(enumerate(ddpm.scheduler.timesteps)):
+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) because I am most familar
-        # with what these images look like
+        # 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
 
-plt.imshow(x[0].cpu().detach().numpy())
-plt.show()
+# Display image
+plt.imshow(np.transpose(x[0].cpu().detach().numpy(), (1,2,0)) + 0.5)
 ```
 
-## Example Outputs
+# 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
 
 <img width="608" alt="image" src="https://cdn-uploads.huggingface.co/production/uploads/6349716695ab8cce385f450e/uxDp-0svPAp2dCmTK36rf.png">
 
-## Training Loss
+# 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.
 
 <img width="608" alt="image" src="https://cdn-uploads.huggingface.co/production/uploads/6349716695ab8cce385f450e/XEZb34rdFYaeFckDMyCYm.png">