[Paper Review] PyTorch Connectomics: A Scalable and Flexible Segmentation Framework for EM Connectomics

Zudi Lin, Donglai Wei+2 more|arXiv (Cornell University)|Dec 10, 2021

Cell Image Analysis Techniques27 citations

TL;DR

PyTorch Connectomics (PyTC) is an open-source framework for semantic and instance segmentation of volumetric EM images, enabling scalable data loading, flexible model architectures, multi-target losses, and semi-/self-supervised learning, with state-of-the-art CREMI synaptic cleft results.

ABSTRACT

We present PyTorch Connectomics (PyTC), an open-source deep-learning framework for the semantic and instance segmentation of volumetric microscopy images, built upon PyTorch. We demonstrate the effectiveness of PyTC in the field of connectomics, which aims to segment and reconstruct neurons, synapses, and other organelles like mitochondria at nanometer resolution for understanding neuronal communication, metabolism, and development in animal brains. PyTC is a scalable and flexible toolbox that tackles datasets at different scales and supports multi-task and semi-supervised learning to better exploit expensive expert annotations and the vast amount of unlabeled data during training. Those functionalities can be easily realized in PyTC by changing the configuration options without coding and adapted to other 2D and 3D segmentation tasks for different tissues and imaging modalities. Quantitatively, our framework achieves the best performance in the CREMI challenge for synaptic cleft segmentation (outperforms existing best result by relatively 6.1$\%$) and competitive performance on mitochondria and neuronal nuclei segmentation. Code and tutorials are publicly available at https://connectomics.readthedocs.io.

Motivation & Objective

• Address the need for a scalable, flexible segmentation toolkit tailored to EM connectomics datasets of varying sizes and modalities.
• Provide a modular framework enabling semantic and instance segmentation across neurons, synapses, mitochondria, and nuclei.
• Support multi-task learning, semi-/self-supervised learning, and configurable architectures to utilize limited labels and large unlabeled data.
• Offer easy customization through configuration files to handle anisotropic 2D/3D data and diverse imaging modalities.

Proposed method

• Modular four-component system: data module for scalable data loading (Volume Dataset and Tile Dataset), model module for 2D/3D or mixed architectures with multiple targets, engine module for distributed/non-distributed training and multi-GPU inference, and configuration module (YACS) for reproducibility.
• Hybrid-representation learning that allows learning multiple targets with multiple losses (e.g., foreground mask, instance contour, signed distance transform) from a single shared network.
• Active and semi-supervised learning workflows (two-stream active query; self-training with pseudo-labels) to leverage unlabeled data.
• Flexible network architectures enabling anisotropic data handling through staged 2D then 3D convolutions, with options for U-Net, DeepLab-like backbones, and mixed 2D/3D convolutions.
• Data augmentations including non-spatial (contrast, brightness, gamma, missing-part) and spatial (rescale, misalignment) and test-time augmentations (up to 16 variants).
• Distributed training with DataParallel or DistributedDataParallel, non-overlapping data partitioning to balance foreground densities, and an inference workflow that chunks large volumes for parallel processing.

Experimental results

Research questions

• RQ1 How can a single framework support both semantic and instance segmentation tasks in EM connectomics with configurable multi-target losses?
• RQ2 Can scalable data loading and distributed training enable practical segmentation of petabyte-scale EM datasets?
• RQ3 Do hybrid-representation and semi-/self-supervised learning improve segmentation performance when labeled data are scarce?
• RQ4 How effective are anisotropic 2D/3D architectures and mixed convolutions for nanometer-resolution EM data?
• RQ5 What are the gains of active learning and self-training on challenging connectomics segmentation tasks?

Key findings

• PyTC achieves state-of-the-art CREMI synaptic cleft segmentation, with semi-supervised learning improving the overall CREMI score by 6.1% over prior best results.
• On synaptic cleft data, the SSL approach ranks first among CREMI submissions and outperforms existing methods on ADGT, ADF, and CREMI metrics.
• For mitochondria semantic segmentation, PyTC-based models surpass prior methods and even approach or exceed human annotation performance with simple post-processing.
• For mitochondria instance segmentation on MitoEM, the U3D-BC/U3D-BCD models achieve competitive AP-75 scores and rank third among challenge participants, highlighting effectiveness in large-scale instance segmentation.
• On neuronal nuclei, the U3D-BC and U3D-BCD variants outperform several baselines (Cellpose, Stardist), with the distance-map variant delivering notable gains (approx. 22% improvement in overall AP).
• Across datasets, PyTC demonstrates scalable data handling (Volume vs Tile datasets), multi-GPU training/inference, and flexible multi-target learning that benefits low-label regimes.