htem / raygun

An image enhancement and segmentation pipeline for generating connectomic reconstructions from X-ray holographic nanotomography, using CycleGANs, Local Shape Descriptors, and Mutex Watershed. Built with PyTorch, Daisy, and Gunpowder.

Raygun

The goal of this repository is to provide an extendable toolbox for large-scale experimentation with deep learning techniques for image enhancement and segmentation of N-dimensional biological imaging data. It is designed to support high-performance computing clusters and utilize GPU-acceleration.

Training datasets, as well as model checkpoints used in the paper will be made accessible soon!

Install:

Run the following -->

conda create -n raygun python=3.9 tensorflow pytorch torchvision torchaudio cudatoolkit=11.3 affogato -c pytorch -c nvidia -c conda-forge 
conda activate raygun
pip install git+https://github.com/htem/raygun

Should you run into gcc / boost errors when conda/pip installing raygun, try installing libboost first:

sudo apt-get update
sudo apt-get install libboost-all-dev

Example train:

raygun-train path/to/train_config.json

raygun's Issues

Fix padding on JAX ResNet

Torch/JAX don't share many parallels when it comes to padding in Haiku - reconfiguring how to implement padding for JAX resents

Fix JAX resource exhaustion/multi-threading

JAX network testing running into two main issues, look to collapse one at a time:

• Single GPU training is able to load the untrained network onto the GPU but unable to allocate resources to run convolutions
• Multi-threaded training splits the first input dimension (batch size=16) into separate (4,4) dimensions, causing the rank of the input tensors to be expanded

Determine Requirements

Determine install requirements and include in setup.py and consider autoupdate from requirements.txt?

raygun/setup.py

Line 10 in 31932c5

requires=[] #TODO

Split up utils.py into the different network files

Add Tensorboard logging to JAX

This may be helpful.

Add Docstrings

(Once refactoring complete and fully tested.)
Add Docstrings to all functions. Priority order:

• class.__init__'s
• standalone functions
• other major functions

Add testing modules to use with pytest

• Adapt simple denoising training to test network architectures (See freezeNorm.py)
• Determine other tests (and add to this list)

JAX refactor for networks

Refactor UNet, Residual UNet, MIRNet2D, and Utils from PyTorch to JAX:

• UNet
• Residual UNet
• MIRNet2D
• Utils

Inheritance args bug with ResNet & NLayerDiscriminator

__init__ functions claim to be overloaded with parameters upon calling, though initialization seems correct:

• ResNet call for ResNetBlock2D/3D calls:
  TypeError: __init__() takes from 2 to 3 positional arguments but 6 were given

• NLayerDiscriminator call for each dimension class:
  TypeError: __init__() takes from 2 to 3 positional arguments but 13 were given