This repository contains the official Pytorch implementation of our paper: EMDiffuse: A Diffusion-based Deep Learning Method Augmenting Ultrastructural Imaging and Volume Electron Microscopy accepted by Nature Communications.

EMDiffuse offers a toolkit for applying diffusion models to electron microscopy (EM) images, designed to enhance ultrastructural imaging in EM and extend the capabilities of volume electron microscopy (vEM). We have tailored the diffusion model for EM applications, developing EMDiffuse-n for EM denoising, EMDiffuse-r for EM super-resolution, and vEMDiffuse-i and vEMDiffuse-a for generating isotropic resolution data from anisotropic volumes in vEM.

A selection of model weights is available at EMDiffuse_model_weight. Download them and place them in the ./experiments folder. The vEMDiffuse-i model was trained on the Openorgnelle liver dataset. And vEMDiffuse-a was trained on the MICrONS multi-area dataset.

All results, including training and inference, will be stored in a newly created folder under ./experiments.

Running the diffusion process on a GPU is highly recommended for both training and testing.

For more information, please visit our webpage: https://www.haibojianglab.com/emdiffuse.

Should you have any questions regarding the code, please do not hesitate to contact us.

1. Experiments on anisotropic volumes with fewer layers (e.g., 128 layers) are now possible. Given our input size 256, we duplicate the anisotropic volume along the axial axis before proceeding with vEMDiffuse-a training. (15/08/2023)

Please install PyTorch (=1.13) before you run the code. We strongly recommend you install Anaconda3, where we use Python 3.8.

conda create --name emdiffuse python=3.8
conda activate emdiffuse
pip install -r requirements.txt

Explore our Jupyter notebooks for step-by-step tutorials:

• 2D EM denoising
• 2D EM super-resolution
• Isotropic vEM reconstruction with isotropic training data
• Isotropic vEM reconstruction without isotropic training data

In order to run the notebooks, install jupyter in your conda environment or use Google Colab.

pip install jupyter

Download the dataset from https://zenodo.org/records/10205819.

Register and crop patches for model training. For our dataset with multiple noise levels:

cd RAFT/core
python register.py --path /data/EMDiffuse_dataset/brain_train --tissue Brain --patch_size 256 --overlap 0.125

Replace the path with your dataset's file path. patch_size should be a power of two or divisible by 8, and overlap sets the overlap ratio of adjacent patches.

For transfer learning on other samples, replace tissue with the target, such as the Liver, Heart, or BM.

For your own denoise dataset with file structure:

Denoise_dataset
	img
		1.tif
		2.tif
		3.tif
		...
	gt
		1.tif
		2.tif
		3.tif
		...

cd RAFT/core
python register_custom.py --path /data/EMDiffuse_dataset/brain_train --patch_size 256 --overlap 0.125

Replace the path with your dataset's file path.

cd ../..
python run.py -c config/EMDiffuse-n.json -b 16 --gpu 0,1,2,3 --port 20022 --path /data/EMDiffuse_dataset/brain_train/denoise/train_wf --lr 5e-5

gpu denotes the GPU devices to be used during training. Multiple GPU training is supported.

Both the model's state and its training metrics are automatically saved within a newly created directory, ./experiments/train_EMDiffuse-r_time. Here, time is a placeholder for the actual timestamp when the training session begins.

Download the dataset from https://zenodo.org/records/10205819.

python test_pre.py --path /data/EMDiffuse_dataset/brain_test --task denoise

Download the model weight and place them in the ./experiments folder.

python run.py -p test -c config/EMDiffuse-n.json --gpu 0 -b 60 --path /data/EMDiffuse_dataset/brain_test/denoise_test_crop_patches --resume ./experiments/EMDiffuse-n/best --mean 1 --step 1000

The diffusion model samples one plausible solution from the learned solution distribution. mean denotes the number of outputs to generate and averaging (each output and averaged output will be saved). resume indicates the path to the model's weight file. step controls the number of diffusion steps, with more steps generally leading to higher image quality.

Download the dataset from https://zenodo.org/records/10205819.

cd RAFT/core
python super_res_register.py --path /data/EMDiffuse_dataset/brain_train --patch_size 128 --overlap 0.125

cd ../..
python run.py -c config/EMDiffuse-r.json -b 16 --gpu 0,1,2,3 --port 20022 --path /data/EMDiffuse_dataset/brain_train/zoom/train_wf --lr 5e-5

Download the dataset from https://zenodo.org/records/10205819.

python test_pre.py --path /data/EMDiffuse_dataset/brain_test --task super

Download the model weight and place them in the ./experiments folder.

python run.py -p test -c config/EMDiffuse-r.json --gpu 0 -b 60 --path /data/EMDiffuse_dataset/brain_test/super_test_crop_patches --resume ./experiments/EMDiffuse-r/best --mean 1 --step 1000

These instructions are tailored for our denoising and super-resolution datasets. For your own dataset, you may need to adjust the cropping and registration codes to suit your data format. Here's a simple demonstration for performing inference on a dataset with the following file structure:

test_images:
		image1.tif
		image2.tif
		....

python crop_single_file.py --path ./test_images --task denoise 

python run.py -p test -c config/EMDiffuse-n.json --gpu 0 -b 60 --path ./test_images/denoise_test_crop_patches/ --resume ./experiments/EMDiffuse-n/best --mean 1 --step 1000

Download or prepare your vEM training data. The training file structure should be like:

vEM_data
	0.tif // The first layer
	1.tif // The second layer
	...
	n.tif // The (n+1)th layer

python run.py -c config/vEMDiffuse-i.json -b 16 --gpu 0,1,2,3 --port 20022 --path ./vEM_data -z 6 --lr 5e-5

-z means the subsampling factor of the Z axis. For example, to reconstruct an 8 nm x 8 nm x 8 nm volume from an 8 nm x 8 nm x 48 nm volume, the subsampling factor should be 6.

Slice along YZ view:

python vEMa_pre.py --path ./vEM_data

Training

python run.py -c config/vEMDiffuse-a.json -b 16 --gpu 0,1,2,3 --port 20022 --path ./vEM_data/transposed -z 6 --lr 5e-5

To test, prepare an anisotropic volume. Execute the model using the appropriate configuration and model weights for isotropic reconstruction.

vEM_test_data
	0.tif // The first layer
	1.tif // The second layer
	...
	n.tif // The (n+1)th layer

python run.py -p test -c config/vEMDiffuse-i.json --gpu 0 -b 16 --path ./vEM_test_data/ -z 6 --resume ./experiments/vEMDiffuse-i/best --mean 1 --step 200

Adjust the model weight directory to where your best model weights are saved.