DOCSEARCHbyAlgolia
Submit Docs

MedSegDiff PyTorch

repository·main·Indexed 42 days ago

https://github.com/lucidrains/med-seg-diff-pytorch

A PyTorch implementation of the MedSegDiff paper for medical image segmentation. It utilizes Diffusion Probabilistic Models (DDPM) with Fourier space feature filtering and enhanced feature-level conditioning to generate segmentation masks.

Tokens
830
Snippets
3
Records
3
Agent score
4%

What's inside MedSegDiff

  2. Use MedSegDiff for training and sampling

    main

    To use MedSegDiff, you need to initialize a Unet and wrap it with the MedSegDiff class.

    Training Workflow

    During training, you pass both the segmented_imgs (the target masks) and the input_imgs (the source images) to the diffusion model to compute the loss.

    Sampling Workflow

    After training, use the .sample() method to generate predicted segmentation masks from unsegmented input images.

    Note: Input images should be normalized from 0 to 1.

    import torch
from med_seg_diff_pytorch import Unet, MedSegDiff

# Initialize the UNet backbone
model = Unet(
    dim = 64,
    image_size = 128,
    mask_channels = 1,          # segmentation has 1 channel
    input_img_channels = 3,     # input images have 3 channels
    dim_mults = (1, 2, 4, 8)
)

# Wrap with MedSegDiff
diffusion = MedSegDiff(
    model,
    timesteps = 1000
).cuda()

# --- Training Step ---
segmented_imgs = torch.rand(8, 1, 128, 128)  # target masks (normalized 0 to 1)
input_imgs = torch.rand(8, 3, 128, 128)       # source images

loss = diffusion(segmented_imgs, input_imgs)
loss.backward()

# --- Inference/Sampling ---
# after a lot of training
pred = diffusion.sample(input_imgs)     # pass in your unsegmented images
print(pred.shape)                       # predicted segmented images - (8, 1, 128, 128)

  3. Train MedSegDiff using accelerate

    main

    Run the training script using accelerate launch.

    Available CLI Arguments

    • --mask_channels: Number of channels in the segmentation mask.
    • --input_img_channels: Number of channels in the input image.
    • --image_size: Spatial resolution of the images.
    • --data_path: Path to the dataset.
    • --dim: Dimension for the UNet.
    • --epochs: Number of training epochs.
    • --batch_size: Batch size.
    • --scale_lr: Scale the learning rate.
    • --gradient_accumulation_steps: Number of steps for gradient accumulation.
    • --self_condition: Enables self-conditioning (conditioning with the mask generated so far).
    accelerate launch driver.py --mask_channels=1 --input_img_channels=3 --image_size=64 --data_path='./data' --dim=64 --epochs=100 --batch_size=1 --scale_lr --gradient_accumulation_steps=4