cyto_dl.models.im2im.diffusion_autoencoder module#

class cyto_dl.models.im2im.diffusion_autoencoder.DiffusionAutoEncoder(*args, **kwargs)[source]#

Bases: BaseModel

[DiffusionAutoencoder](https://arxiv.org/abs/2111.15640) for representation learning. Code is based on the [MONAI generative tutorial](https://github.com/Project-MONAI/GenerativeModels/blob/main/tutorials/generative/2d_diffusion_autoencoder/2d_diffusion_autoencoder_tutorial.ipynb)

Parameters:

  • autoencoder (nn.Module) – model network to denoise the diffusion image (conditioned on the latent generated by the semantic encoder)

  • spatial_inferer (Inferer) – Inferer to use for splitting large images into patches during inference

  • image_shape (Sequence[int]) – C[Z]YX shape of the input images

  • condition_key (str) – key to access condition images in batch

  • noise_scheduler (Scheduler) – beta noise scheduler

  • diffusion_inferer (Inferer) – Inferer to use for diffusion sampling

  • loss (nn.Module) – loss function to use for training. Should have no reduction.

  • semantic_encoder (nn.Module) – model network to encode the condition image

  • diffusion_key (Optional[str]) – key to access diffusion images in batch. If None, defaults to condition_key

  • n_inference_steps (int) – number of noise steps used during inference. Must be less than the number of train steps used in your noise scheduler, and can be much fewer due to DDIM sampling

  • save_dir=”./” – directory to save images during training and validation

  • save_images_every_n_epochs (int) – Image saving frequency

  • n_noise_samples (Optional[int]) – Number of noise samples to average for latent walk

  • train_encoder (bool) – Whether to train the semantic encoder

  • gamma (float) – Minimum SNR for loss weighting. If negative, no weighting is applied

  • noise_schedule (str) – beta noise schedule. Options are ‘inverse_cosine’ from the SODA paper or see the MONAI docs for other options

  • **base_kwargs – Additional arguments passed to BaseModel

configure_optimizers()[source]#
encode_image(x)[source]#
forward(x_cond, x_diff)[source]#
generate_from_latent(cond: Tensor, save_name: str = 'generated_image', n_noise_samples: int | None = None, average: bool = True, save: bool = True, batch_size: int = 3)[source]#

Generate images from latent features. If average is True, average over n_noise_samples, otherwise make a composite.

Parameters:#

cond: torch.Tensor

latent features to condition the diffusion model

save_name: str

name to save the generated image

n_noise_samples: int

number of noise samples to generate

average: bool

Whether to average the generated images. If False, composite the images.

save: bool

Whether to save the generated image. If False, return the generated image.

batch_size: int

batch size for generating images

model_step(stage, batch, batch_idx)[source]#
predict_step(batch, batch_idx)[source]#
save_example(stage, cond_img, diff_img)[source]#

Save the sequence of denoising steps.

validation_step(batch, batch_idx)[source]#
cyto_dl.models.im2im.diffusion_autoencoder.inverted_cosine_beta_schedule(num_train_timesteps, s=0.008)[source]#

inverted cosine schedule as proposed in https://arxiv.org/pdf/2311.17901.pdf