Package: lsds
Local Shape Descriptors (for Neuron Segmentation)This repository contains code to compute Local Shape Descriptors (LSDs) from an instance segmentation. LSDs can then be used during training as an auxiliary target, which we found to improve boundary prediction and therefore segmentation quality. Read more about it in ourpaperand/orblog post.PaperBlog PostQuick 2d ExamplesNotebooksExample networks & pipelinesParallel processingCite:@article{sheridan_local_2021,title={Local Shape Descriptors for Neuron Segmentation},url={https://www.biorxiv.org/content/10.1101/2021.01.18.427039v1},urldate={2021-01-20},journal={bioRxiv},author={Sheridan, Arlo and Nguyen, Tri and Deb, Diptodip and Lee, Wei-Chung Allen and Saalfeld, Stephan and Turaga, Srinivas and Manor, Uri and Funke, Jan},year={2021}}Notes:Tested on Ubuntu 18.04 with Python 3.This is not production level software and was developed in a pure research environment. Therefore some scripts may not work out of the box. For example, all paper networks were originally written using now deprecated tensorflow/cudnn versions and rely on an outdated singularity container. Because of this, the singularity image will not build from the current recipe - if replicating with the current implementations, please reach out for the singularity container (it is too large to upload here). Alternatively, consider reimplementing networks in pytorch (recommended - seeTraining).Post-proccesing steps were designed for use with a specific cluster and will need to be tweaked for individual use cases. If the need / use increases then we will look into refactoring, packaging and distributing.Currently, several post-processing scripts (e.gwatershed) are located inside this repo which creates more dependencies than needed for using the lsds. One forseeable issue is that agglomeration requires networkx==2.2 for the MergeTree and boost is required forfunlib.segment. We have restructured the repo to uselsd.trainandlsd.postsubmodules. For just calculating the lsds, it is sufficient to uselsd.train, e.g:fromlsd.trainimportlocal_shape_descriptorQuick 2d ExamplesThe following tutorial allows you to run in the browser using google colab. In order to replicate the tutorial locally, create a conda environment and install the relevant packages. E.g:conda create -n lsd_test python=3conda activate lsd_testpip install matplotlib scikit-image gunpowderpip install git+https://github.com/funkelab/lsd.gittutorial:NotebooksExamble colab notebooks are locatedhere. You can download or run below (control + click open in colab). When running a notebook, you will probably get the message: "Warning: This notebook was not authored by Google". This can be ignored, you can run anyway.We uploaded ~1.7 tb of data (raw/labels/masks/rags etc.) to an s3 bucket. The following tutorial shows some examples for accessing and visualizing the data.Data download:If implementing the LSDs in your own training pipeline (i.e pure pytorch/tensorflow), calculate the LSDs on a label array of unique objects and use them as the target for your network (see quick 2d examples above for calculating).The following tutorials show how to set up 2D training/prediction pipelines usingGunpowder. It is recommended to follow them in order (skip the basic tutorial if familiar with gunpowder).Note:Google Colab can sometimes be slow especially due to data I/O. These notebooks will run much faster in a jupyter notebook on a local gpu, but the Colab versions should provide a starting point.Basic Gunpowder tutorial:Train Affinities:Train LSDs:Train MTLSD:Inference (using pretrained MTLSD checkpoint):Watershed, agglomeration, segmentation:Bonus notebooks:Training using sparse ground truth (useful if you only have a subset of training data but still want dense predictions):Ignore regions during training (useful if you want the network to learn to predict zeros in certain regions, eg glia ids):Train lsds on non-em data with pytorch:Example networks & pipelinesThere are some example networks and training/prediction pipelines from the fib25 datasethere.TrainingSince networks in this paper were implemented in Tensorflow, there was a two step process for training. First the networks were created using themknet.pyfiles. This saved tensor placeholders and meta data in config files that were then used for both training and prediction. The mknet files used the now deprecated mala repository to create the networks. If reimplementing in Tensorflow, consider migrating tofunlib.learn.tensorflow.If using Pytorch, the networks can just be created directly inside the train scripts since placeholders aren't required. For example, the logic from this tensorflowmknet scriptand this tensorflowtrain scriptcan be condensed to this pytorchtrain script.For training an autocontext network (e.gacrlsd), the current implementation learns the LSDs in afirst pass. A saved checkpoint is then used when creating thesecond passin order topredict LSDsprior to learning the Affinities. One could modify this to use a single setup and remove the need for writing the LSDs to disk.InferenceBy default, the predict scripts (example) contain the worker logic to be distributed by the scheduler during parallel processing (see below).If you just need to process a relatively small volume, it is sometimes not necessary to use blockwise processing. In this case, it is recommended to use ascan node, and specify input/output shapes + context. An example can be found in the inference colab notebook above.Similar to training, the current autocontext implementations assume the predicted LSDs are written to a zarr/n5 container and then used as input to the second pass to predict affinities. This can also be changed to predict on the fly if needed.Visualizations of example training/prediction pipelinesVanilla affinitiestraining:AutocontextLSDandaffinitiesprediction:Parallel processingIf you are running on small data then this section may be irrelevant. See theWatershed, agglomeration, segmentationnotebook above if you just want to get a sense of obtaining a segmentation from affinities.Example processing scripts can be foundhereWe create segmentations following the approach inthis paper. Generally speaking, after training a network there are five steps to obtain a segmentation:Predict boundaries (this can involve the use of LSDs as an auxiliary task)Generate supervoxels (fragments) using seeded watershed. The fragment centers of mass are stored as region adjacency graph nodes.Generate edges between nodes using hierarchical agglomeration. The edges are weighted by the underlying affinities. Edges with lower scores are merged earlier.Cut the graph at a predefined threshold and relabel connected components. Store the node - component lookup tables.Use the lookup tables to relabel supervoxels and generate a segmentation.Everything was done in parallel using daisy (github,docs), but one could use multiprocessing or dask instead.For our experiments we usedMongoDBfor all storage (block checks, rags, scores, etc) due to the size of the data. Depending on use case, it might be better to read/write to file rather than mongo. See watershed for further info.The following examples were written for use with the Janelia LSF cluster and are just meant to be used as a guide. Users will likely need to customize for their own specs (for example if using a SLURM cluster).Need to installfunlib.segmentandfunlib.evaluateif using/adapting segmentation/evaluation scripts.InferenceThe worker logic is located in individualpredict.pyscripts (example). Themaster scriptdistributes usingdaisy.run_blockwise. The only need for MongoDb here is for the block check function (to check which blocks have successfully completed). To remove the need for mongo, one could remove the check function (remember to also removeblock_done_callbackinpredict.py) or replace with custom function (e.g check chunk completion directly in output container).Example roi config{"container":"hemi_roi_1.zarr","offset":[140800,205120,198400],"size":[3000,3000,3000]}Example predict config{"base_dir":"/path/to/base/directory","experiment":"hemi","setup":"setup01","iteration":400000,"raw_file":"predict_roi.json","raw_dataset":"volumes/raw","out_base":"output","file_name":"foo.zarr","num_workers":5,"db_host":"mongodb client","db_name":"foo","queue":"gpu_rtx","singularity_image":"/path/to/singularity/image"}WatershedThe worker logic is located in a singlescriptwhich is then distributed by themaster script. By default the nodes are stored in mongo using aMongoDbGraphProvider. To write to file (i.e compressed numpy arrays), you can use theFileGraphProviderinstead (inside the worker script).Example watershed config{"experiment":"hemi","setup":"setup01","iteration":400000,"affs_file":"foo.zarr","affs_dataset":"/volumes/affs","fragments_file":"foo.zarr","fragments_dataset":"/volumes/fragments","block_size":[1000,1000,1000],"context":[248,248,248],"db_host":"mongodb client","db_name":"foo","num_workers":6,"fragments_in_xy":false,"epsilon_agglomerate":0,"queue":"local"}AgglomerateSame as watershed.Worker script,master script. Change to FileGraphProvider if needed.Example agglomerate config{"experiment":"hemi","setup":"setup01","iteration":400000,"affs_file":"foo.zarr","affs_dataset":"/volumes/affs","fragments_file":"foo.zarr","fragments_dataset":"/volumes/fragments","block_size":[1000,1000,1000],"context":[248,248,248],"db_host":"mongodb client","db_name":"foo","num_workers":4,"queue":"local","merge_function":"hist_quant_75"}Find segmentsIn contrast to the above three methods, whencreating LUTsthere just needs to be enough RAM to hold the RAG in memory. The only thing done in parallel is reading the graph (graph_provider.read_blockwise()). It could be adapted to use multiprocessing/dask for distributing the connected components for each threshold, but if the rag is too large there will be pickling errors when passing the nodes/edges. Daisy doesn't need to be used for scheduling here since nothing is written to containers.Example find segments config{"db_host":"mongodb client","db_name":"foo","fragments_file":"foo.zarr","edges_collection":"edges_hist_quant_75","thresholds_minmax":[0,1],"thresholds_step":0.02,"block_size":[1000,1000,1000],"num_workers":5,"fragments_dataset":"/volumes/fragments","run_type":"test"}Extract segmentationThisscriptdoes use daisy to write the segmentation to file, but doesn't necessarily require bsub/sbatch to distribute (you can run locally).Example extract segmentation config{"fragments_file":"foo.zarr","fragments_dataset":"/volumes/fragments","edges_collection":"edges_hist_quant_75","threshold":0.4,"block_size":[1000,1000,1000],"out_file":"foo.zarr","out_dataset":"volumes/segmentation_40","num_workers":3,"run_type":"test"}Evaluate volumesEvaluateVoi scores. Assumes dense voxel ground truth (not skeletons). This also assumes the ground truth (and segmentation) can fit into memory, which was fine for hemi and fib25 volumes assuming ~750 GB of RAM. The script should probably be refactored to run blockwise.Example evaluate volumes config{"experiment":"hemi","setup":"setup01","iteration":400000,"gt_file":"hemi_roi_1.zarr","gt_dataset":"volumes/labels/neuron_ids","fragments_file":"foo.zarr","fragments_dataset":"/volumes/fragments","db_host":"mongodb client","rag_db_name":"foo","edges_collection":"edges_hist_quant_75","scores_db_name":"scores","thresholds_minmax":[0,1],"thresholds_step":0.02,"num_workers":4,"method":"vanilla","run_type":"test"}Evaluate annotationsFor the zebrafinch, ground truth skeletons were used due to the size of the dataset. These skeletons were cropped, masked, and relabelled for the sub Rois that were tested in the paper. Weevaluatedvoi, erl, and the mincut metric on the consolidated skeletons. The current implementation could be refactored / made more modular. It also usesnode_collectionswhich are now deprecated in daisy. To use with the current implementation, you should checkout daisy commit39723ca.Example evaluate annotations config{"experiment":"zebrafinch","setup":"setup01","iteration":400000,"config_slab":"mtlsd","fragments_file":"foo.zarr","fragments_dataset":"/volumes/fragments","edges_db_host":"mongodb client","edges_db_name":"foo","edges_collection":"edges_hist_quant_75","scores_db_name":"scores","annotations_db_host":"mongo client","annotations_db_name":"foo","annotations_skeletons_collection_name":"zebrafinch","node_components":"zebrafinch_components","node_mask":"zebrafinch_mask","roi_offset":[50800,43200,44100],"roi_shape":[10800,10800,10800],"thresholds_minmax":[0.5,1],"thresholds_step":1,"run_type":"11_micron_roi_masked"}