Supervoxel Agglomeration using Graph Neural Networks

Implementation based on pytorch_geometric.

Abstract

Fully automating neural circuit reconstruction from electron microscopy (EM) data has the potential to provide a huge leap for neuroscience. The final step of many neuron reconstruction pipelines consists of agglomerating small compact fragments to form entire neurons. This work explores using Graph Neural Networks to learn that task from the Region Adjacency Graph (RAG) of said fragments, combined with features learned by a Convolutional Neural Network directly from EM data. For this purpose we present a generalizable, geometric Graph Neural Network called RagNet. It labels RAG edges as merge or split, which allows the direct extraction of segmentations. In contrast to previous agglomeration methods, the RagNet allows an increase in the context considered for making an edge merge decision. This benefit is empirically confirmed by applying RagNet on top of an edge-wise prediction method and increasing the class-balanced accuracy by ten percentage points. Lastly, we fuse our implementation with an existing large-scale neuron reconstruction pipeline and report initial results on a recently imaged Drosophila melanogaster brain dataset.

Installation

1. Set up a python3 conda environment, then pip install -r requirements.txt or conda env create -f environment.yml
2. Conda install pytorch as specified here
3. Install pytorch_geometric with ./install_pytorch_geometric.sh, or one-by-one as specified here
4. Adapt gnn_agglomeration/config.py and run main.py

Experiment tracking with Sacred

1. To install MongoDB on MacOS, execute brew install mongodb, then start it as a service with brew services start mongodb
2. Access the mongodb via the mongo shell with mongo
3. In there, set up a new database called sacred with use sacred
4. Insert a dummy entry with db.dummy.insert({"dummy":"dummy"})
5. To set up Omniboard, follow the steps here
6. Execute omniboard -m <host>:27017:sacred to start up Omniboard at localhost:9000