kmouts/mednist-hand-ddpm-2d

Denoising Diffusion Probabilistic Models with MedNIST Dataset

Model Information

This repository contains a Denoising Diffusion Probabilistic Model (DDPM) built using MONAI. The model is designed to generate synthetic 2D images, specifically focusing on medical images from the MedNIST dataset.

Training Data Details

The model was trained on the MedNIST dataset, specifically utilizing images from the 'Hand' class. The following transformations were applied to the training data:

• LoadImaged: Loads medical images from files.
• EnsureChannelFirstd: Ensures images have a channel-first dimension.
• ScaleIntensityRanged: Scales image intensities from [0, 255] to [0, 1].
• RandAffined: Applies random affine transformations (rotation, translation, scaling) for data augmentation.

Usage Instructions

After training, the model can be used to generate new synthetic images. By providing a random noise input, the inferer.sample method can iteratively denoise it to produce a realistic image corresponding to the training data distribution.

Technical Specifications

Framework

The model is built using the MONAI framework, leveraging PyTorch for deep learning operations.

Architecture

The core of the diffusion model is a DiffusionModelUNet configured as follows:

• spatial_dims=2: For 2D image processing.
• in_channels=1: Single input channel (grayscale images).
• out_channels=1: Single output channel.
• num_channels=(128, 256, 256): Number of channels in the UNet's encoder and decoder blocks.
• attention_levels=(False, True, True): Attention mechanisms are applied at the second and third levels of the UNet.
• num_res_blocks=1: One residual block per level.
• num_head_channels=256: Number of head channels for attention. The noise scheduling is handled by a DDPMScheduler with num_train_timesteps=1000.

Hyperparameters

The model was trained with the following key hyperparameters:

• Learning Rate (lr): 2.5e-5
• Batch Size: 128
• Number of Epochs: 67

Training Duration

The total training time for the model was approximately 5224.60 seconds.

Dependencies

The following Python packages are essential for running this project:

• monai-weekly[tqdm]
• monai-generative
• matplotlib
• numpy
• torch

Citations

[1] Ho et al. "Denoising Diffusion Probabilistic Models" (https://arxiv.org/abs/2006.11239)