Reworked Implementation of RLN

Original Paper:

Incorporating the image formation process into deep learning improves network performance.

This repository has three main branches:

Requirements

• gmerlin access or an nvidia gpu with cuda & drivers installed.
• apptainer

Containers

This project uses apptainer as the main container engine. Basic commands are:

apptainer build target.sif definiton.def #build the definition.def into target.sif
apptainer exec targe.sif {Your Command} #run {Your Command} in the container.
apptainer run target.sif #run the preconfigured command of the container.

See documentation for more details. There are three containers currently provided:

• basic_image: This is the last available official build of tensorflow 1.14 without gpu support.
• basic_image_gpu: This is the last available official build of tensorflow 1.14 with gpu support.
• nvidia_image_gpu: This is the latest build of the nvidia tensorflow 1.x project. Provides modern gpu support for tensorflow 1.15

On GPU Acceleration

The container basic_image_gpu is compiled for cuda 10. This is a very old version and does not support modern GPUs. This triggers an on-the-fly recompilation cascade that can take hours depending on how new the GPU is (newer means more stuff to recompile).

The nvidia container solves this problem. This container is provided by Nvidia and contains TF 1 that was compiled for a newer version of CUDA. It has the downside of being ca. 8 GB without the data. Still, this is the only way to get the code running on modern gpus without triggering hours long recompilation every time

Dataset & Model

There is a test dataset and weights available from the authors in the respective folders. There are also matlab scripts provided by the authors that generate synthetic data. See the folder for more details.

Status

We have tested that the main script RLN_single.py is working as intended for the modes:

• TR: Training
• TS: Inference

Other modes and other scripts might have unknown bugs.

The matlab scripts in Phantom_generate work and can be used to generate images blurred using custom PSFs. We only tested this for the provided PSF.tif.

Usage

To use the RLn network on your data you need to train and adapt the model to the noise characteristics of your microscope. This can be done by the following steps:

1. Determine the PSF of the microscope.
2. Use the Matlab scripts in Phantom_generate to generate phantoms that were blured using the PSF of your microscope.
3. Train the network using the synthetic data generated by the Matlab scripts.
4. Use the network on real world images.

Training

Data Structure

Before training, we need to create our folder structure. This can be on merlin in the /data directory or locally if you are using a local machine

{path to your data}
├── logs
├── test
│   ├── ground_truth
│   ├── input
│   └── output_rl
└── train
    ├── ground_truth
    ├── input
    ├── model_rl
    └── output_rl

The ground truth and blurred versions of the images should have the same names but be in different folders.

Synthetic Data

The matlab scripts in Phantom_generate output the images in the following format :

DATA
├── ground_truth
│   └── 1.tif
├── input_noise
│   └── 1.tif
└── input_no_noise
    └── 1.tif

The number of images depends on a variable in Phantom_generate.m. We suggest 200 images. This will generate 200 separate .tif files for each output category. For the training we used the output in "input_noise". These have an added gaussian noise on top of the PSF. In our experience this helps stop the network from overfitting.

Running the scripts

After creating the files as in Section "Data Structure", you should copy DATA/ground_truth/* to {path to your data}/train/ground_truth and DATA/input_noise/ to {path to your data}/train/input/. Then you can use:

apptainer run --nv --bind {path to your data}:/data {container name}.sif --mode TR

to train the network.

The --bind option makes your {path to your data} accessible in the container as /data. The scripts are programmed to look at this directory as their base directory. The --nv option tells apptainer to make the nvidia driver and cuda binaries available to the container. The container automatically runs the scripts. The arguments after [...].sif are passed on to the python scripts. Running it with --help will show the available settings. These settings correspond to setting variables in the original paper with default values taken from the paper.

Monitoring the Progress

Using tensorboard it is possible to monitor the training progress in real-time from a web browser. Since the code is in legacy Tensorflow you also need to use the legacy version of TensorBoard which is available in the provided containers. A basic workflow could be:

1. Run the training on merlin.

   apptainer run --nv --bind {path to your data}:/data nvidia_image_gpu:.sif --mode TR

2. Mount {path to your data}/logs in merlin to a local older using sshfs:

   sshfs user@merlin-l-001.psi.ch:{path to your data}/logs mnt-merlin

3. Bind mnt-merlin to the container and open a shell in the container:

   apptainer --bind mnt-merlin:/mnt-merlin shell basic_image.sif

4. Run tensorboard inside the container

   tensorbaord --logdir /mnt-merlin   

5. Now you can access the website from a browser on your computer.

Running Inference

Put your tif files in {path to you data}/test/input. Run:

    apptainer run --nv --bind {path to your data}:/data {container name}.sif --mode TS

The output will be in test/output_rl. This command will use the latest checkpoint in train/model_rl. The checkpoint is automatically saved there during training.