Flood-Filling Networks (FFNs) are a class of neural networks designed for instance segmentation of complex and large shapes, particularly in volume EM datasets of brain tissue.

For more details, see the related publications:

• https://arxiv.org/abs/1611.00421
• https://doi.org/10.1101/200675

This is not an official Google product.

You can create the Conda environment required to run this code by using the following command to launch a job on the Ibex cluster via Slurm.

sbatch bin/create-conda-env.sbatch

We provide a sample coordinate file for the FIB-25 validation1 volume included in third_party. Due to its size, that file is hosted in Google Cloud Storage. Google Cloud SDK, which is already installed in the Conda environment, can be used to download the data files. If you haven't used it before, you will need to set it up by running the following command.

gcloud auth application-default login

You will be given a link to copy and paste into your browser which will then given you a code that you will need to copy and paste back into terminal when prompted.

Once initialized, you can create a local copy of the required data files with the following command.

gsutil rsync -r gs://ffn-flyem-fib25/ third_party/neuroproof_examples

The entire data preparation, training, and inference pipeline can be launched on the Ibex cluster via Slurm with the following command.

./bin/launch-pipeline.sh

FFN networks can be trained with the train.py script, which expects a TFRecord file of coordinates at which to sample data from input volumes.

There are two scripts to generate training coordinate files for a labeled dataset stored in HDF5 files: compute_partitions.py and build_coordinates.py.

compute_partitions.py transforms the label volume into an intermediate volume where the value of every voxel A corresponds to the quantized fraction of voxels labeled identically to A within a subvolume of radius lom_radius centered at A. lom_radius should normally be set to (fov_size // 2) + deltas (where fov_size and deltas are FFN model settings). Every such quantized fraction is called a partition. Sample invocation:

python compute_partitions.py \
    --input_volume third_party/neuroproof_examples/validation_sample/groundtruth.h5:stack \
    --output_volume third_party/neuroproof_examples/validation_sample/af.h5:af \
    --thresholds 0.025,0.05,0.075,0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9 \
    --lom_radius 24,24,24 \
    --min_size 10000

This stage of the pipline can be launched individually using the following command.

sbatch bin/compute-partitions.sbatch

build_coordinates.py uses the partition volume from the previous step to produce a TFRecord file of coordinates in which every partition is represented approximately equally frequently. Sample invocation:

python build_coordinates.py \
     --partition_volumes validation1:third_party/neuroproof_examples/validation_sample/af.h5:af \
     --coordinate_output third_party/neuroproof_examples/validation_sample/tf_record_file \
     --margin 24,24,24

This stage of the pipline can be launched individually using the following command.

sbatch bin/build-coordinates.sbatch

Once the coordinate files are ready, you can start training the FFN with:

python train.py \
    --train_coords gs://ffn-flyem-fib25/validation_sample/fib_flyem_validation1_label_lom24_24_24_part14_wbbox_coords-*-of-00025.gz \
    --data_volumes validation1:third_party/neuroproof_examples/validation_sample/grayscale_maps.h5:raw \
    --label_volumes validation1:third_party/neuroproof_examples/validation_sample/groundtruth.h5:stack \
    --model_name convstack_3d.ConvStack3DFFNModel \
    --model_args "{\"depth\": 12, \"fov_size\": [33, 33, 33], \"deltas\": [8, 8, 8]}" \
    --image_mean 128 \
    --image_stddev 33

This stage of the pipline can be launched individually using the following command.

sbatch bin/train.sbatch

Note that both training and inference with the provided model are computationally expensive processes. We recommend a GPU-equipped machine for best results, particularly when using the FFN interactively in a Jupyter notebook. Training the FFN as configured above requires a GPU with 12 GB of RAM. You can reduce the batch size, model depth, fov_size, or number of features in the convolutional layers to reduce the memory usage.

The training script is not configured for multi-GPU or distributed training. For instructions on how to set this up, see the documentation on Distributed TensorFlow.

We provide two examples of how to run inference with a trained FFN model. For a non-interactive setting, you can use the run_inference.py script:

  python run_inference.py \
    --inference_request="$(cat configs/inference_training_sample2.pbtxt)" \
    --bounding_box 'start { x:0 y:0 z:0 } size { x:250 y:250 z:250 }'

which will segment the training_sample2 volume and save the results in the results/fib25/training2 directory.

This stage of the pipline can be launched individually using the following command.

sbatch bin/run-inference.sbatch

Two files will be produced: seg-0_0_0.npz and seg-0_0_0.prob. Both are in the npz format and contain a segmentation map and quantized probability maps, respectively. In Python, you can load the segmentation as follows:

from ffn.inference import storage
seg, _ = storage.load_segmentation('results/fib25/training2', (0, 0, 0))

We provide sample segmentation results in results/fib25/sample-training2.npz. For the training2 volume, segmentation takes ~7 min with a P100 GPU.

For an interactive setting, check out ffn_inference_demo.ipynb. This Jupyter notebook shows how to segment a single object with an explicitly defined seed and visualize the results while inference is running.

Both examples are configured to use a 3d convstack FFN model trained on the validation1 volume of the FIB-25 dataset from the FlyEM project at Janelia.

Please see doc/manual.md.