Skip to content

Commit

Permalink
Add ZMD (#2)
Browse files Browse the repository at this point in the history
* Add ZMD

* Add ImageNet data sample

* Add a Dockerfile and other improvements

* Update README

---------

Co-authored-by: Maria Wyrzykowska <[email protected]>
  • Loading branch information
Ruruthia and Maria Wyrzykowska authored Sep 27, 2024
1 parent 0cd1553 commit 729ddce
Show file tree
Hide file tree
Showing 136 changed files with 536 additions and 186 deletions.
5 changes: 4 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -42,5 +42,8 @@ venv/
.vscode/

# Data & models
/data/
data/
models/

# Exclude 'imnet_sample' from being ignored
!data/imnet_sample/
55 changes: 55 additions & 0 deletions Dockerfile
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
# Use an official lightweight image as a base
FROM ubuntu:22.04

# Set environment variables to avoid interactive prompts during installation
ENV DEBIAN_FRONTEND=noninteractive

# Install system dependencies
RUN apt-get update && apt-get install -y \
curl \
wget \
bzip2 \
ca-certificates \
libglib2.0-0 \
libxext6 \
libsm6 \
libxrender1 \
git \
&& apt-get clean \
&& rm -rf /var/lib/apt/lists/*

# Install micromamba
RUN curl -L https://micro.mamba.pm/api/micromamba/linux-64/latest | tar -xvj -C /usr/local/bin --strip-components=1 bin/micromamba

# Set up working directory
WORKDIR /app

# Copy current directory contents into the container
COPY . .

# Set environment name
ENV ENV_NAME=cvdm

ENV MAMBA_ROOT_PREFIX=/opt/micromamba
ENV PATH=$MAMBA_ROOT_PREFIX/envs/$ENV_NAME/bin:$PATH
ENV PYTHONPATH=$MAMBA_ROOT_PREFIX/envs/$ENV_NAME/lib/python3.10/site-packages:$PYTHONPATH

RUN apt-get update && apt-get install ffmpeg libsm6 libxext6 -y
# Set up micromamba environment, install dependencies and pip packages
RUN micromamba create -n $ENV_NAME -y && \
micromamba install -n $ENV_NAME \
tensorflow-gpu==2.15.0 \
keras==2.15.* \
matplotlib==3.8.0 \
tqdm==4.65.0 \
scikit-learn==1.4.2 \
scikit-image==0.22.0 \
einops==0.7.0 \
neptune==1.10.2 -y && \
/usr/local/bin/micromamba run -n $ENV_NAME pip3 install opencv-python==4.9.0.80 \
tensorflow-addons==0.23.0 \
cupy-cuda12x==13.3.0 && \
/usr/local/bin/micromamba run -n $ENV_NAME python -m pip install .

# Default command when the container starts (optional)
CMD ["/bin/bash"]
33 changes: 22 additions & 11 deletions README.md
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# Conditional Variational Diffusion Models

This code implements the Conditional Variational Diffusion Models as described [in the paper](https://arxiv.org/abs/2312.02246).
Diffusion models have become popular for their ability to solve complex problems where hidden information needs to be estimated from observed data. Among others, their use is popular in image generation tasks. These models rely on a key hyperparameter of the variance schedule that impacts how well they learn, but recent work shows that allowing the model to automatically learn this hyperparameter can improve both performance and efficiency. Our CVDM package implements Conditional Variational Diffusion Models (CVDM) as described [in the paper](https://arxiv.org/abs/2312.02246) that build on this idea, with the addition of [Zero-Mean Diffusion (ZMD)](https://arxiv.org/pdf/2406.04388), a technique that enhances performance in certain imaging tasks, aiming to make these approaches more accessible to researchers.

## Where to get the data?

Expand All @@ -14,34 +14,45 @@ It is assumed that for:
- BioSR phase task, data can be found in the directory specified as dataset_path in configs/biosr_phase.yaml, in one file, y.npy (ground truth). Input to the model will be generated based on the ground truth.
- ImageNet super-resolution task, data can be found in the directory specified as dataset_path in configs/imagenet_sr.yaml as a collection of JPEG files. Input to the model will be generated based on the ground truth.
- ImageNet phase task, data can be found in the directory specified as dataset_path in configs/imagenet_phase.yaml as a collection of JPEG files. Input to the model will be generated based on the ground truth.
- HCOCO phase evaluation task, data can be found in the directory specified as dataset_path in configs/hcoco_phase_eval.yaml as a collection of JPEG files. Input to the model will be generated based on the ground truth.
- HCOCO phase evaluation task, data can be found in the directory specified as dataset_path in configs/hcoco_phase.yaml as a collection of JPEG files. Input to the model will be generated based on the ground truth.

## How to prepare environment?

Run the following code:

We provide a Dockerfile to prepare the environment. Run the following code in the root of this repository:
```
docker build -t my-image .
docker run -it my-image
```
Inside the image run:
```
conda create -n myenv python=3.10
conda activate myenv
pip install -r requirements.txt
pip install -e .
eval "$(micromamba shell hook --shell bash)"
micromamba activate cvdm
```

If you encounter issues with cupy installation (required only for the phase tasks) such as [these](https://github.com/cupy/cupy/issues/8466), you can modify the `cvdm/utils/phase_utils.py` to use pure numpy.

## How to run the training code?

1. Download the data.
1. Modify the config in `configs/` directory with the path to the data you want to use and the directory for outputs.
2. Run the code from the root directory: `python scripts/train.py --config-path $PATH_TO_CONFIG --neptune-token $NEPTUNE_TOKEN`.
1. Download the data or use the sample data available in the data/ directory. The sample data is a fraction of the ImageNet dataset and can be used with configs `imagenet_sr_sample.yaml` or `imagenet_phase_sample.yaml`. You can also use your own data as long as it is in ".npy" format. To do so, use the task type "other".
2. Modify the config in `configs/` directory with the path to the data you want to use and the directory for outputs. For the description of each parameter, check the documentation in `cvdm/configs/` files.
3. Run the code from the root directory: `python scripts/train.py --config-path $PATH_TO_CONFIG --neptune-token $NEPTUNE_TOKEN`.

`--neptune-token` argument is optional.

## How to run the training code?

## How to run the evaluation code?

1. Download the data.
1. Modify the config in `configs/` directory with the path to the data you want to use and the directory for outputs.
2. Run the code from the root directory: `python scripts/eval.py --config-path $PATH_TO_CONFIG --neptune-token $NEPTUNE_TOKEN`.

`--neptune-token` argument is optional.

## How to contribute?

To contribute to the software or seek support, please leave an issue or pull request.

## License
This repository is released under the MIT License (refer to the LICENSE file for details).

7 changes: 5 additions & 2 deletions configs/biosr.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,18 @@ model:
noise_model_type: "unet"
alpha: 0.001
load_weights: null
load_mu_weights: null
snr_expansion_n: 1
zmd: False
diff_inp: False

training:
lr: 0.0001
epochs: 10

eval:
output_path: "outputs/biosr"
generation_timesteps: 1000
generation_timesteps: 200
checkpoint_freq: 1000
log_freq: 10
image_freq: 100
Expand All @@ -26,5 +29,5 @@ data:
im_size: 256

neptune:
name: "Virtual_Stain"
name: "CVDM"
project: "mlid/test"
7 changes: 5 additions & 2 deletions configs/biosr_phase.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,18 @@ model:
noise_model_type: "unet"
alpha: 0.001
load_weights: null
load_mu_weights: null
snr_expansion_n: 1
zmd: False
diff_inp: True

training:
lr: 0.0001
epochs: 10

eval:
output_path: "outputs/biosr"
generation_timesteps: 1000
generation_timesteps: 200
checkpoint_freq: 1000
log_freq: 10
image_freq: 100
Expand All @@ -26,5 +29,5 @@ data:
im_size: 256

neptune:
name: "Virtual_Stain"
name: "CVDM"
project: "mlid/test"
9 changes: 6 additions & 3 deletions configs/hcoco_phase_eval.yaml → configs/hcoco_phase.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,18 @@ model:
noise_model_type: "unet"
alpha: 0.001
load_weights: null
load_mu_weights: null
snr_expansion_n: 1
zmd: False
diff_inp: False

training:
lr: 0.0001
epochs: 100

eval:
output_path: "outputs/hcoco"
generation_timesteps: 1000
generation_timesteps: 200
checkpoint_freq: 1000
log_freq: 10
image_freq: 100
Expand All @@ -22,9 +25,9 @@ eval:
data:
dataset_path: "/bigdata/casus/MLID/maria/hcoco_sample"
n_samples: 100
batch_size: 1
batch_size: 2
im_size: 256

neptune:
name: "Virtual_Stain"
name: "CVDM"
project: "mlid/test"
11 changes: 7 additions & 4 deletions configs/imagenet_phase.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -4,27 +4,30 @@ model:
noise_model_type: "unet"
alpha: 0.001
load_weights: null
load_mu_weights: null
snr_expansion_n: 1
zmd: False
diff_inp: False

training:
lr: 0.0001
epochs: 100

eval:
output_path: "outputs/imagenet"
generation_timesteps: 1000
generation_timesteps: 200
checkpoint_freq: 1000
log_freq: 10
image_freq: 100
val_freq: 200
val_len: 100

data:
dataset_path: "/bigdata/casus/MLID/maria/imagenet_sample"
dataset_path: "/bigdata/imnet"
n_samples: 100
batch_size: 1
batch_size: 2
im_size: 256

neptune:
name: "Virtual_Stain"
name: "CVDM"
project: "mlid/test"
33 changes: 33 additions & 0 deletions configs/imagenet_phase_sample.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
task: "imagenet_phase"

model:
noise_model_type: "unet"
alpha: 0.001
load_weights: null
load_mu_weights: null
snr_expansion_n: 1
zmd: False
diff_inp: False

training:
lr: 0.0001
epochs: 100

eval:
output_path: "outputs/imagenet"
generation_timesteps: 200
checkpoint_freq: 1000
log_freq: 10
image_freq: 100
val_freq: 200
val_len: 10

data:
dataset_path: "data/imnet_sample"
n_samples: 100
batch_size: 2
im_size: 256

neptune:
name: "CVDM"
project: "mlid/test"
9 changes: 6 additions & 3 deletions configs/imagenet_sr.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -3,27 +3,30 @@ model:
noise_model_type: "unet"
alpha: 0.001
load_weights: null
load_mu_weights: null
snr_expansion_n: 1
zmd: False
diff_inp: True

training:
lr: 0.0001
epochs: 100

eval:
output_path: "outputs/imagenet"
generation_timesteps: 1000
generation_timesteps: 200
checkpoint_freq: 1000
log_freq: 10
image_freq: 100
val_freq: 200
val_len: 100

data:
dataset_path: "/bigdata/casus/MLID/maria/imagenet_sample"
dataset_path: "/bigdata/imnet"
n_samples: 100
batch_size: 2
im_size: 256

neptune:
name: "Virtual_Stain"
name: "CVDM"
project: "mlid/test"
32 changes: 32 additions & 0 deletions configs/imagenet_sr_sample.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
task: "imagenet_sr"
model:
noise_model_type: "unet"
alpha: 0.001
load_weights: null
load_mu_weights: null
snr_expansion_n: 1
zmd: False
diff_inp: True

training:
lr: 0.0001
epochs: 100

eval:
output_path: "outputs/imagenet"
generation_timesteps: 200
checkpoint_freq: 1000
log_freq: 10
image_freq: 1000
val_freq: 2000
val_len: 10

data:
dataset_path: "data/imnet_sample"
n_samples: 100
batch_size: 2
im_size: 256

neptune:
name: "CVDM"
project: "mlid/test"
4 changes: 0 additions & 4 deletions cvdm/architectures/components/residual_block.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,6 @@
from tensorflow.keras.layers import Concatenate, Conv2D, Dropout
from tensorflow_addons.layers import GroupNormalization

from cvdm.architectures.components.conditional_instance_normalization import (
ConditionalInstanceNormalization,
)


def resblock(
x: tf.Tensor, noise_embedding: tf.Tensor, n_out_channels: int, dropout: float = 0.0
Expand Down
10 changes: 6 additions & 4 deletions cvdm/architectures/sr3.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,17 @@
from typing import Tuple

import tensorflow as tf
from tensorflow.keras.activations import swish
from tensorflow.keras.layers import AveragePooling2D, Conv2D, Input, UpSampling2D
from tensorflow.keras.models import Model
from tensorflow_addons.layers import GroupNormalization

from cvdm.architectures.components.attention_block import attention_block
from cvdm.architectures.components.deep_residual_block import (
deep_resblock,
up_deep_resblock,
)
from cvdm.architectures.components.residual_block import resblock, up_resblock
from tensorflow.keras.models import Model
from tensorflow.keras.activations import swish
from tensorflow.keras.layers import AveragePooling2D, Conv2D, Input, UpSampling2D
from tensorflow_addons.layers import GroupNormalization


def upsample(x: tf.Tensor, use_conv: bool = False) -> tf.Tensor:
Expand Down
Loading

0 comments on commit 729ddce

Please sign in to comment.