Pixel Diffusion UNet – LoDoChallenge (DM4CT)

This repository contains the pretrained pixel-space diffusion UNet presented in the paper DM4CT: Benchmarking Diffusion Models for Computed Tomography Reconstruction (ICLR 2026).


πŸ”¬ Model Overview

This model learns a prior over CT reconstruction images using a denoising diffusion probabilistic model (DDPM). It operates directly in pixel space (not latent space).

  • Architecture: 2D UNet (Diffusers UNet2DModel)
  • Input resolution: 512 Γ— 512
  • Channels: 1 (grayscale CT slice)
  • Training objective: Ξ΅-prediction (standard DDPM formulation)
  • Noise schedule: Linear beta schedule
  • Training dataset: Low Dose Grand Challenge (LoDoChallenge)
  • Intensity normalization: Rescaled to (-1, 1)

This model is intended to be combined with data-consistency correction for CT reconstruction tasks as detailed in the DM4CT benchmark.


πŸ“Š Dataset: Low Dose Grand Challenge

Source: AAPM Low Dose CT Grand Challenge

Preprocessing steps:

  • Train/test split.
  • Rescale reconstructed slices to (-1, 1).
  • No geometry information is embedded in the model.

The model learns an unconditional image prior over medical CT slices.


🧠 Training Details

  • Optimizer: AdamW
  • Learning rate: 1e-4
  • Hardware: NVIDIA A100 GPU
  • Training script: train_pixel.py

πŸš€ Usage

You can use this model with the diffusers library:

from diffusers import DDPMPipeline

# Load the pretrained pipeline
pipeline = DDPMPipeline.from_pretrained("jiayangshi/lodochallenge_pixel_diffusion")

# Generate a sample (unconditional CT slice prior)
image = pipeline().images[0]
image.save("generated_ct_slice.png")

Citation

@inproceedings{
  shi2026dmct,
  title={{DM}4{CT}: Benchmarking Diffusion Models for Computed Tomography Reconstruction},
  author={Shi, Jiayang and Pelt, Dani{\"e}l M and Batenburg, K Joost},
  booktitle={The Fourteenth International Conference on Learning Representations},
  year={2026},
  url={https://openreview.net/forum?id=YE5scJekg5}
}
Downloads last month
21
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support

Paper for jiayangshi/lodochallenge_pixel_diffusion