Large-scale image segmentation based on distributed clustering algorithms

Many approaches to 3D image segmentation are based on hierarchical clustering of supervoxels into image regions. Here we describe a distributed algorithm capable of handling a tremendous number of supervoxels. The algorithm works recursively, the regions are divided into chunks that are processed independently in parallel by multiple workers. At each round of the recursive procedure, the chunk size in all dimensions are doubled until a single chunk encompasses the entire image. The final result is provably independent of the chunking scheme, and the same as if the entire image were processed without division into chunks. This is nontrivial because a pair of adjacent regions is scored by some statistical property (e.g. mean or median) of the affinities at the interface, and the interface may extend over arbitrarily many chunks. The trick is to delay merge decisions for regions that touch chunk boundaries, and only complete them in a later round after the regions are fully contained within a chunk. We demonstrate the algorithm by clustering an affinity graph with over 1.5 trillion edges between 135 billion supervoxels derived from a 3D electron microscopic brain image.


Large-scale image segmentation based on distributed clustering algorithms Ran Lu, Aleksandar Zlateski and H. Sebastian Seung
Abstract-Many approaches to 3D image segmentation are based on hierarchical clustering of supervoxels into image regions. Here we describe a distributed algorithm capable of handling a tremendous number of supervoxels. The algorithm works recursively, the regions are divided into chunks that are processed independently in parallel by multiple workers. At each round of the recursive procedure, the chunk size in all dimensions are doubled until a single chunk encompasses the entire image. The final result is provably independent of the chunking scheme, and the same as if the entire image were processed without division into chunks. This is nontrivial because a pair of adjacent regions is scored by some statistical property (e.g. mean or median) of the affinities at the interface, and the interface may extend over arbitrarily many chunks. The trick is to delay merge decisions for regions that touch chunk boundaries, and only complete them in a later round after the regions are fully contained within a chunk. We demonstrate the algorithm by clustering an affinity graph with over 1.5 trillion edges between 135 billion supervoxels derived from a 3D electron microscopic brain image.

I. INTRODUCTION
Reconstructing neural connectivity from electron microscopic (EM) images has become an important method for neuroscience. In light of advancing technologies, the size and quality of image stacks collected have been growing exponentially in recent years. High quality segmentations of these images are essential to extract precise information of neural circuits. However, neurons have complex morphology and segmenting them from grey scale EM images is a challenging task.
Convolutional networks are now standard for EM image segmentation [Lee et al., 2019b]. Here we focus on one of these approaches [Dorkenwald et al., 2019]: First we use convolutional networks to calculate inter-voxel affinities between nearest neighbors. The affinities are then converted to segmentation by classical algorithms including watershed and agglomeration. Variants of this approach include adding affinities between more distant voxels as an auxiliary task during training [Lee et al., 2019a], improved training techniques [Funke et al., 2019] etc.
After the affinities are determined, the image is segmented using classical algorithms that do not depend on machine learning. The convolutional networks only predict the affinities within their fields of view (FOV). For large objects that cannot be included in the FOV, watershed and agglomeration algorithm extend the affinities into the entire image stack to create complete segments. If the models' predictions were perfect, this could be done by calculating the connected components over the affinities, however in reality the errors and noises in the output of the convolutional networks are unavoidable. Using sophisticated segmentation algorithms [Wolf et al., 2018, Beier et al., 2017 helps achieve higher quality by suppressing these errors through statistical means. For this purpose, most of the popular agglomeration algorithms take an iterative approach, agglomerate fragments in descending order of the affinities between them, and update the affinities through statistical methods after each agglomeration. This introduces a global dependency on the input data. A local change in the input can cascade through the iteration process and induce errors in a far region of the output. Thus, it is common to process the entire input together. Currently, most of the proposed methods were demonstrated on small image stacks that completely fit into the working memory of a single machine. When segmenting larger images, the input images are divided into (overlapping) chunks and processed independently; After all chunks are segmented they are stitched together by heuristics [Funke et al., 2019]. As aforementioned algorithms have global dependencies, this perchunk processing could introduce additional errors.
Here, we propose a novel approach to segmenting large image stacks. The input images are similarly divided into chunks so they can be distributed in computer clusters or cloud computing services and processed in parallel. However, the algorithm we use to agglomerate within each chunk and stitch chunks together guarantees that no errors are introduced throughout the process. Our method generalizes to a large class of hierarchical clustering algorithms -as long as affinities between clusters are updated by linkage criteria that satisfy the reducibility condition [Müllner, 2011]. Using notations we will explain in more details in the following sections, considering three clusters I, J and K, this is true for any linkage criterion satisfying: where A (I, J) is the affinity between the clusters I and J. Many popular linkage criteria used in EM image segmentation fall into this category. This is not a surprise, since the input of these clustering algorithms are affinities predicted by convolutional networks. These affinities can be interpreted as estimated probabilities of two regions belonging to a single segment; updating them by some statistical averaging is natural. It is well known in the literature, when the linkage criteria satisfy the reducibility condition, some agglomeration operations can be performed out of order. In this paper we point out the reducibility condition also allow algorithms to perform agglomeration within partial input and avoid introducing errors due to the incomplete local information. Our algorithm treats the supervoxels touching the chunk boundary in a special manner; we use this information to avoid fictitious mergers or splits. After we merge all the clusters that can be agglomerated, we are left with an intermediate, partial, segmentation containing the residual clusters that cannot be safely agglomerated within the chunk, as some non-local information might affect the decision of whether to merge two clusters or not. We suggest an elegant solution to merge the partially segmented chunks: we simply collect the residual clusters and apply the same algorithm on them, now with no boundaries. More generally we can start from small chunks, and gradually merge them into bigger and bigger chunks, during which we applying the same agglomeration algorithm at each step, and updating the boundaries accordingly. Once we merge all chunks together, we obtain a segmentation of the entire input image stack. Using our method, all the segments are agglomerated by the same algorithm. There is no need to develop a separate stitching method. It not only simplifies the segmentation pipeline, also eliminates errors introduced by the stitching. In fact, it can be proven that the segmentation obtained this way is independent of the chunking scheme, thus identical to the segmentation created by a single machine that has all the global information.

AGGLOMERATIVE CLUSTERING ALGORITHMS
Image segmentation classifies voxels into domains representing different objects in the input image, it can be treated as a clustering problem. The input is an affinity graph G 0 = (V 0 , E 0 , A 0 ). The nodes in V 0 represent single voxel clusters, E 0 ⊂ V 0 × V 0 contains edges between the nodes, and the map A 0 : E 0 → R stores the affinities associated with the edges. For image segmentation, there is a natural definition of locality. Following this intuition, we often only consider the correlations between direct neighboring voxels. This makes the affinity graph sparsely connected. For example if we only consider correlations between voxels and their direct neighbors in of a 3D image stack, |E 0 | ≈ 3|V 0 |. Locality is a crucial assumption to our distributed algorithm: if voxels can have arbitrary correlations at any distance, we must consider the entire G 0 when clustering voxels. This assumption can be extended to include long range correlations studied in [Lee et al., 2017, Wolf et al., 2018, as long as the long range correlations are still confined within some bounded FOV. Throughout the agglomeration, we track the process with a region adjacency graph (RAG) G = (V, E, A). At the beginning the RAG is simiply the affinity graph G = G 0 . As we agglomerate the voxels together, we update G by merging nodes and updating the weights between them according to the linkage criterion.
In the end we collect the edges corresponding to the clustering operations to construct the dendrogram.
Various linkage criteria have been invented for clustering problems in different fields. For image segmentation, only consider affinities between direct neighboring voxels, we can use the affinities between the boundary voxels to determine the affinities between two clusters. 1 . In this case linkage criteria are often some statistical property of the boundary voxels' affinity distribution: max, mean, quantiles, etc. This can be extended to incorporate extra input based on the clusters' intrinsic properties. In [Zlateski and Seung, 2015] the max affinity linkage criterion was augmented by size thresholds. One can also include prior knowledge of the underlying image stack to improve the segmentation. For example, assigning negativing affinities to voxel pairs with different semantic labels to prevent false mergers or reweighting the affinities to account for varying image quality.
There are a few more considerations specific to image segmentation we want to mention. First, our goal is to cluster voxels into domains representing the objects we want to study. There is little interest to agglomerate everything into a single cluster and obtain a complete dendrogram. For image segmentation, it is sufficient to stop clustering at a threshold T and output the truncated dendrogram or a flat segmentation. Also, as we alluded in the Introduction section, commonly we start clustering from supervoxels 2 . Starting from single voxels can be considered as a special case. These supervoxels can be created by faster but less robust algorithms like watershed or obtained from an existing (over)segmentation. The affinities between supervoxels are readily defined by the linkage criterion. Using supervoxels can reduce the size of the input RAG and simplify the computation. It also allows us to perform multi-stage agglomeration by switching algorithm and parameters, treating the segments of the previous step as input supervoxels to optimize the final results. Finally, since there are no intrinsic differences, we will use clusters, supervoxels, segments as interchangeable terms.
III. CHUNK-BASED AGGLOMERATIVE CLUSTERING ALGORITHM To explain our chunk-based clustering algorithm, we start by reviewing a generic algorithm similar to the one studied in [Kurita, 1991]. The input data is a RAG G (V, E, A). We also supply a threshold T and stop the agglomeration after the affinities of all edges are lower than it. Before starting the agglomeration, we create a heap H from the edge list E, which are used to extract the edge with the highest affinity to perform the next clustering operations. The output of the agglomeration is a dendrogram D The global nature of Algorithm 1 is obvious in line 5: We extract the edge with the highest affinity from the heap and agglomerate the two nodes connected by it. If the input RAG is incomplete, we will likely agglomerate a different set of edges, and the output dendrogram D will be different. 1 If we are interested in long range affinities we have to include the long range edges across the boundary as well 2 Small clusters of voxels that are likely belong to the same object   [Murtagh, 1985, Müllner, 2011, Bruynooghe, 1977, de Rham, 1980. Consider three segments I, J and K, and the affinities between them are represented as A (I, J) , A (I, K) and A (J, K), reducibility condition requires: Intuitively, this condition indicates there is no reversal. The affinity between the merged cluster I ∪ J and K cannot be higher than both A (I, K) and A (J, K).
We call a segment J the nearest neighbor of another segment I, if the affinity A (I, J) satisfy: Unless completely isolated, there is a nearest neighbor for each segment. Segment I and J are called mutual nearest neighbors when I is also J's nearest neighbor. We call the edge between I and J is a local maximal edge. One can readily see when the linkage criterion satisfy the reducibility condition, mutual nearest neighbors always merge with one another. Agglomerating the rest of the RAG cannot create an edge with higher affinity associated with I or J because the reducibility condition forbids such reversal. Similarly, agglomerating I and J together will not create edges with affinities higher than A (I, J). We can agglomerate I and J out of the standard order, the original dendrogram can be restore by a simple permutation. One can extend this reasoning to a rigorous proof using induction [Müllner, 2011]. Reducibility condition is famous for allowing O(|V | 2 ) clustering algorithms like the nearest-neighbor chain [Benzécri, 1982, Juan, 1982, which have been studied extensively in the literature. We point out that for image segmentation, combining with the locality assumption, reducibility conditiona allow us to identify and agglomerate local maximal edges with partial inputs. The modified algorithm suited for chunked input based on this idea is presented in Algorithm 2.
Algorithm 2 Clustering algorithm with chunked input 1: procedure AGGLOMERATE CHUNK(G C , B C , T ) 2: H ← Priority queue from G C 3: In algorithm 2, we change our notations of the input and output from G, D to G C , D C to emphasis we intend to apply it to a chunk. To make the new algorithm aware of the chunked input, we include an extra input B C . It contains all the boundary supervoxels touching the artificial chunk boundaries 3 . These supervoxels may be split among chunks, and the edges associated with them in G C may be incomplete. Because of these boundary supervoxels, some of the segments cannot be agglomerated within the chunk. We call these segments "frozen segments" and keep track of them in a set F . At the beginning, F = B C . The edge {u, v} we extracted 3 not to be confused with the dataset's real boundaries 4 in line 7 by definition is the edge with the highest affinity in G C . if neither u nor v is in F , u and v are mutual nearest neighbors and we can agglomerate them immediately. Otherwise it means whether we can agglomerate u and v depends on agglomeration decisions involving segments in F . We cannot merge u and v before we resolve those pending agglomerations. We add u and v to F to forbid agglomeration involving u and v in future iterations 4 . The frozen edges are saved into a graph G C F . Later we process them correctly after we stitch the chunks together.
The rest of the changes are straightforward: we can no longer discard all the edges below the agglomeration threshold T , because some of them are frozen edges belonging to G C F . In the end, we return the partial dendrogram, and the frozen edges for later steps. We draw a red line between a segment and its nearest neighbor, with an arrow pointing from it to its nearest neighbor. For mutual nearest neighbors, we draw a red line with arrows in both directions. The numbers on top of the lines are the affinities between the segments. When we follow the arrows, we will eventually find a pair of mutual nearest neighbors, and the segments we visited are called a nearest-neighbor chain. In this graph, segment A and D, segment G and H are the mutual nearest neighbors. It is straightforward to verify that we can agglomerate G and H first without altering the dendrogram.
IV. DISTRIBUTED CLUSTERING ALGORITHM Algorithm 2 is the corner stone of our distributed clustering algorithm. It not only can agglomerate chunked input, but also can stitch those chunked results together to create segmentation of the entire dataset. For each chunk we get a partial dendrogram D C , and a fronzen graph G C F . For simplicity we assume the chunks are not overlapping. In this situation, the segments across the chunk boundaries are avoided by Algorithm 2, stitching them into a intermediate segmentation is trivial. We can merge all the G C F together into a "residual" graph G res representing the RAG of the intermediate segmentation by matching the segments and merging the incomplete edges. The edges remain in G res are edges cannot be safely agglomerated within individual chunks. Now the RAG are extened to the entire dataset in G res , we can apply algorithm 2 again, with B an empty set. The result is a dendrogram D res . Combine it with partial dendrograms D C s we recover the dendrogram of the entire input image stack. Fig. 2. A RAG cut into two chunks. The graph is the same as the one we showed in Figure 1. Here we draw segments in squares and use the contact surfaces between them to represent the edges. The dashed lines indicate the chunk boundaries. The segments A, B, C, and I in blue squares are boundary supervoxels. When we agglomerate segments within each chunk, we can not merge segments A and D, because A is a boundary supervoxels and we cannot see all its neighbors in either chunk. On the other hand, we can merge segments G and H within the right chunk. After we stitch the two chunks, the segments left can be clustered in the right order. The two-step example above can be extended to a generic recursive algorithm 3. The input is a RAG G, and a boundary segment set B. If the chunk is small, we use AGGLOM-ERATE CHUNK described in algorithm 2 to agglomerate it directly. If the input is too big, we call a DIVIDE CHUNK function to cut it into smaller chunks. The DIVIDE CHUNK function returns a collection of subchunks, each including its edge list G C and boundary supervoxels B C . We then pass the subchunks to algorithm 3 again to have them agglomerated. In the end, we collect the frozen edges in each subchunk and use the function COMBINE EDGES to combine the partial edges in the subchunks. This gives us a reduced edge list G res which is much smaller than the original edge list G. We apply AGGLOMERATE CHUNK to finish the agglomeration within the original chunk and return the combined dendrogram and the list of the remaining frozen edges.
The algorithm described in Algorithm 3 allows us to distribute the agglomeration workload to multiple computers. With a shared storage to communicate the inputs and outputs, each AGGLOMERATE CHUNK call can run on a separate computer. We can design the function DIVIDE CHUNK and COMBINE EDGES most convenient for the dataset and cluster. The output does not depend on the detailed implementation of these functions.
There are non-trivial dependencies between the tasks, as one can see the AGGLOMERATE CHUNK call in line 14 must wait until the calls in line 10 have finished. The memory needed to perform the agglomeration task can also vary. It mainly depends on the size of G res , which is determined by the number of frozen edges in the subchunks. Intuitively as we stitch bigger and bigger chunks, the edges we need to process in G res will also increase. There may be also big objects that span a large portion of the dataset. Before the chunks grow large enough to contain entire objects, most of the edges related to them will be frozen, potentially contributing significantly to the memory increase.
As long as we use linkage criteria satisfying the reducibility for G C , B C in C do 10: return D, G F 17: end procedure condition, the generic clustering algorithm 1 and the distributed algorithm 3 are equivalent. We can test linkage criteria with small test inputs using Algorithm 1, and applying it to large datasets using Algorithm 3 V. APPLICATIONS We implemented a system to perform the distributed clustering algorithm we outlined in the previous section. We choose to recursively subdivide the input dataset into eight octants and represent the whole structure as an octree (lower panel of Figure 3). This allows us to identify boundary supervoxels based on their spatial locations. We use Apache Airflow to keep track of the dependencies between the tasks and assign tasks to different types of computers based on the memory requirement. The upper panel of Figure 3 is an example showing how the system processes an image stack. Starting from the EM images, the system generates the affinity maps and supervoxels needed for the agglomeration process. We first agglomerate in small chunks where many supervoxels have to be frozen due to the boundaries. Gradually we stitch the chunks and agglomerate larger and larger objects. The agglomeration process completes when we stitch all the boundaries together. At this point, we only need to agglomerate supervoxels around the remaining boundaries or those frozen due to big objects. Comparing with the naive approach that agglomerates the entire supervoxel input in one pass, the computation resources required are significantly reduced. Algorithm 2 is implemented as part of https://github.com/seung-lab/abiss. And the task orchestration system based on Airflow is available at https: //github.com/seung-lab/seuron. We tested the system with the Phase I dataset of the MI-CrONS project [Turner et al., 2020]. It is a 250×140×90 µm 3 region in a mouse visual cortex imaged at 3.58×3.58×40 nm 3 resolution. The input RAG contains 133 million supervoxels and 1 billion edges between them. We segmented it using Algorithm 1 with the mean affinity linkage criterion. The Starting from an image stack, we first create the nearest neighbor affinity maps using a convolutional neural network. The supervoxel layer is generated by a distributed version of the algorithm described in [Zlateski, 2011]. In the second row, we show how the supervoxels are agglomerated in small chunks and then stitched together. The black lines represent the chunk boundaries. One can clearly see that as the chunks grow, more and more structures are agglomerated, but the supervoxels near the chunk boundaries that may belong to objects across several chunks remain frozen. The panel at the bottom is a schematic diagram of the octree approach. The input image stack is divide into smaller and smaller chunks until they are suitable to be distributed to clusters and cloud computing resources, we then agglomerate the chunks and stitch them back together hierarchically using Algorithm 2. clustering process took 3 days and required 300 GB of memory. After switching to the system based on Algorithm 3, the input was organized into a 7-layer octree to distribute the work to up to 100 n1-highmem-32 instances using Google Compute Engine, The same RAG was agglomerated within 2 hours. The leaf nodes are 512×512×128-voxels chunks. Only 5% of all the merge operations happened in the top node of the octree. The Phase I results of the MICrONS project can be find in https://microns-explorer.org/.
In Phase II of the MICrONS project, we segmented a region roughly 1000 times larger. We managed to run the agglomeration process in two weeks with up to 3000 instances using Google Compute Engine. The input image stack has a bounding box of 1.5 × 1.0 × 0.5 mm 3 with voxel resolution at 8 × 8 × 40 nm 3 . We started from an oversegmentation created by a conservative watershed procedure. The input contained 135 billion supervoxels, the RAG has 1.5 trillion edges. Here we used a 12-layer octree which contained 10 million leaf nodes each represent chunks with up to 256×256×512 voxels.

VI. CONCLUSION AND DISCUSSION
In this paper we present a framework of distributed clustering algorithms and demonstrated the image segmentation system we built based on it. This framework allows us to agglomerate huge image stacks in chunks to utilize clusters or cloud resources to speed up the process, and avoiding introducing artifacts when we stitch the results together. We make a few remarks before concluding the paper: a) There are alternative approaches to segment EM images. For example, Flood filling network (FFN) [Januszewski et al., 2018] played a central role in segmenting a portion of the central brain of the fruit fly [Xu et al., 2020]. In principle one can apply FFN to EM image stack without dividing it into chunks. For the fruit fly, the segmentation are first created in chunks by FFN then merged together conservatively based on the segments in the overlapping regions. While suppressing spurious mergers significantly, the procedure also introduced extra splits. Whether our methods can be adapted to help these approaches is a very interesting open question. b) We started with a generic clustering algorithm (Algorithm 1) and demonstrated how to adapt it into a distributed algorithm when the linkage criterion satisfies the reducibility condition with minimal changes. One can also modify nearestneighbor chain algorithms to consider the chunk boundaries in a similar fashion. In the literature, there are also parallel clustering algorithms based on mutual nearest neighbors that might speed up the process further [Murtagh, 1985]. We tested these implementations in our system, and found they were significant slower than Algorithm 2 when stitching large chunks near the top of the octree. We believe that several factors contribute to this observation. Algorithm 1 has a complexity O (|E| log |E|), while the complexity of a nearest neighbor chain implementation is roughly O (|V |f (V )), here O (f (V )) is the complexity of finding the nearest neighbor for a vertex. Because the RAG is sparse 5 , |E| ∝ |V |. If O (f (V )) is a large constant the nearest neighbor chain algorithm will not be efficient. Due to the way we distributing the tasks, the segments remaining incomplete are often large objects surrounded by many small fragments. Thus for the chunks near the top of the octree, a significant portion of the residual RAG have structures look like Figure 4: Segment A and B are two big structure truncated by the chunk boundary, and the rest of the segments C-Z are small fragments will be absorb into A or B. One can easily see, the nearest neighbor chain construction requires we searching for the nearest neighbor of A or B repeatedly. In our experiment of segmenting large EM image of neural tissue, segments like A and B commonly have several hundreds or even thousands of neighbors. For a brute-force search, O (f (V )) is forbiddingly large. Further more, we notice in Figure 4, segment A, E and B, F cannot be agglomerated in parallel without some careful locking mechanism, otherwise there will be a race condition when they both try to merge the blue edges. c) So far we assumed we divided the input image into non-overlapping chunks. It is also possible to use overlapping chunks in Algorithm 3. The main advantage is that we can cluster more supervoxels in each step: Any supervoxel will have a FOV at least as large as the overlapping region in at least one chunk. Stitching is possible as long as we still using Algorithm 2. The key observation is the following: Suppose chunk C 1 , · · · , C n share an overlapping region which contains Fig. 4. A local structure of RAG common for chunks near the top of the octree. Fragments A and B represents large objects like almost completed neurons. Fragments C to Z are small fragments must be agglomerated to either A or B. a supervoxel v. After clustering each chunk, v now belongs to segment S 1 , · · · , S n . We can prove S 1 , · · · , S n is a totally ordered set when they are ordered by the dendrogram. Based on this statement, to stitch the overlapping chunks, we need to pick the most advanced segment for each supervoxels in the overlapping region. Overlapping chunks can significantly reduce clustering operations pending until the last few steps, making it possible to cluster much larger datasets comparing to the non-overlapping approach. Of course, because we have to agglomerate the overlapping region several times and the overhead can be significant and undesirable for small input. On the left is the stitched segmentation from nonoverlapping 8192×8192×512 chunks. On the right is that from chunks with 2048 × 2048 × 512 overlap. The highlighted regions are where the overlaps located. One can clearly see the fragments are reduced in these regions. The final segmentations are still identical up to a relabeling of the segment ids.