Noise2Noise

To train the Noise2Noise autoencoder on ImageNet, use config.py with the train command.

Default Behavior: By default, this trains a Gaussian denoising network using the ImageNet validation set. On an NVIDIA Titan V GPU, this takes approximately 7.5 hours.

Output: Upon completion, the training process produces a network_final.pickle file in the specific results directory (e.g., results/<run_id>/).

Monitoring: You can inspect the training process using TensorBoard by navigating to the results directory.

# try python config.py train --help for available options
python config.py --desc='-test' train --train-tfrecords=datasets/imagenet_val_raw.tfrecords --long-train=true --noise=gaussian

# To inspect training with TensorBoard:
cd results
tensorboard --logdir .

To generate the TFRecords file required for training the Noise2Noise denoising network using the ImageNet validation set, use the dataset_tool_tf.py script. This process processes approximately 50,000 images and outputs a .tfrecords file.

Command: Run dataset_tool_tf.py with the --input-dir pointing to your ImageNet validation directory and specify the --out path for the resulting TFRecords file.

# This should run through roughly 50K images and output a file called `datasets/imagenet_val_raw.tfrecords`.
python dataset_tool_tf.py --input-dir "<path_to_imagenet>/ILSVRC2012_img_val" --out=datasets/imagenet_val_raw.tfrecords

The training process validates loss against the Kodak Lossless True Color Image Suite. Use the download_kodak.py script to prepare this dataset.

# Download the kodak validation set from http://r0k.us/graphics/kodak/
python download_kodak.py --output-dir=datasets/kodak

Train the MRI denoising network using config_mri.py.

Noise-to-Noise vs Noise-to-Clean: Noise-to-noise training is enabled by default. To switch to noise-to-clean training, edit config_mri.py and set corrupt_targets=False.

Performance: Training for 300 epochs takes approximately 9 hours on an NVIDIA Titan V GPU. The expected average PSNR on the validation set (test_db_clamped) is roughly 31.74 dB.

python config_mri.py

To prepare the IXI-T1 dataset for MRI denoising, use the dataset_tool_mri.py script in three steps:

1. Generate PNGs: Convert the unpacked IXI-T1 dataset into PNG files.
2. Generate Pickles: Convert a subset of those PNGs into training and validation sets (pickled format).

Requirements: You must first download and unpack the IXI-T1 dataset.

# Step 1: Convert IXI-T1 to PNG
# Assumes you have downloaded and untarred IXI-T1 under ~/Downloads/IXI-T1.
python dataset_tool_mri.py genpng --ixi-dir=~/Downloads/IXI-T1 --outdir=datasets/ixi-png

# Step 2: Convert subset to training/validation pickles
python dataset_tool_mri.py genpkl --png-dir=datasets/ixi-png --pkl-dir=datasets

The project is tested with Python 3.6. It is recommended to use Anaconda to manage the environment. Follow these steps to create a clean environment and install the necessary TensorFlow and library dependencies:

1. Create a new conda environment named n2n with Python 3.6.
2. Activate the environment.
3. Install tensorflow-gpu.
4. Upgrade pip.
5. Install dependencies from requirements.txt.

conda create -n n2n python=3.6
conda activate n2n
conda install tensorflow-gpu
python -m pip install --upgrade pip
pip install -r requirements.txt

If ImageNet is unavailable, you can use the BSD300 dataset. After downloading and uncompressing the dataset (e.g., into a ./datasets directory), convert the training images into a TFRecords file using dataset_tool_tf.py.

python dataset_tool_tf.py --input-dir datasets/BSDS300-images/BSDS300/images/train --out=datasets/bsd300.tfrecords