Topology-aware method to segment 3D plant tissues images.
Abstract
The study of genetic and molecular mechanisms underlying tissue morphogenesis has received a lot of attention in biology.
Especially, accurate segmentation of tissues into individual cells plays an important role for quantitative analyzing the development of the growing organs.
However, instance cell segmentation is still a challenging task due to the quality of the image and the ﬁne-scale structure.
Any small leakage in the boundary prediction can merge different cells together, thereby damaging the global structure of the image.
In this paper, we propose an end-to-end topology-aware 3D segmentation method for plant tissues.
The strength of the method is that it takes care of the 3D topology of segmented structures.
The keystone is a training phase and a new topology-aware loss - the CavityLoss - that are able to help the network to focus on the topological errors to ﬁx them during the learning phase.
The evaluation of our method on both ﬁxed and live plant organ datasets shows that our method outperforms state-of-the-art methods (and contrary to state-of-the-art methods, does not require any post-processing stage).
The code of CavityLoss is freely available at https://github.com/onvungocminh/CavityLoss.
Introduction
Learning structural biology from bioimages is a hot topic during the last decade [4, 17, 16, 19].
Especially, segmentation of individual cells is a key step in the study of morphogenesis in a multicellular organism, i.e. counting the cell number of ovule primordia [15], or analysis the asymmetric division of lateral root founder cells [1].
However, volumetric segmentation is still a challenging task and may cause some problems for the downstream analysis of the development of the plants.
In the literature, most of the cell segmentation methods are based on their boundary prediction [15, 19, 20].
Accurate boundary prediction is extremely important because neighboring cells share mutual information on the boundary location.
Thus, even a small leakage on the boundary can merge adjacent regions and damage the global structure.
In light of the very successful trend toward deep learning-based image analysis, various methods have been proposed to leverage the ability of the convolutional neuron networks (CNN) to learn better image representations, thus improving the performance of the boundary prediction [20, 8, 5].
In particular, U-Net [13], which consists of an encoder-decoder architecture, is proved to be useful for biomedical images and has later been further extended to process volumetric data [3].
Although these methods achieve high accuracy, they do not contain any geometrical constraints.
As a consequence, they usually provide discontinuities and broken connections in the segmentation results.
Moreover, due to the limitation of the receptive ﬁeld of the convolutional neural network [21], it can not capture global information of the image.
To address the problem of boundary prediction, several post-processing methods have been proposed to support grouping voxels into cells.
Particularly, the watershed algorithm [2, 10, 12], which is a wellknown method in the ﬁeld of mathematical morphology, has been used to produce the segmentation result from the delineated boundaries.
However, the drawback of the watershed algorithm is when the seed pixels are missing or misclassiﬁed.
In [18], the author proposed a learned watershed algorithm in a reinforcement learning [7] style.
Nevertheless, this method is greedy in terms of training data and computing time.
Another problem with post-processing methods is that they do not integrate into the model as an end-to-end framework, which may slow down the segmentation process.
Recently, several topology-aware segmentation methods have been proposed to preserve the global structure of the image [9, 6, 11].
Commonly, the networks in these methods are trained using a speciﬁc loss - a topological loss - that takes care of the topology of the structures.
A topological loss should satisfy the two following conditions: (1) tolerate with a minor displacement of the boundary; (2) penalize the topological errors.
The idea of topological loss started in [9], then developed in
Topoloss [6], BALoss [11].
However, the detected topological errors from these methods are usually neither relevant nor accurate.
Nevertheless, these methods are mostly designed for 2D images and do not guarantee to preserve the topology.
Especially, a new connectivity-preserving metric clDice [14] has been proposed leveraging the homotopy equivalent between the prediction and its skeleton.
However, clDice treats equally every pixel on the skeleton, thus does not adequately penalize the broken connections.
Therefore, we need a powerful method which is able to maintain the global structure for 3D image segmentation.
In order to tackle the aforementioned topology-preserving problem and provide a better solution to localize the broken connections in object boundaries, we propose a new method based on a 3D network.
To train this network, we propose a new topological loss function called CavityLoss that pays attention to the structure of the image.
Our method identiﬁes topological errors in an efﬁcient and accurate manner while we attempt to reconstruct the ground truth image from the prediction.
Indeed, we leverage the efﬁciency of mathematical morphology in the general and watershed algorithm [2], in particular, to correctly detect the leakages on the 3D boundary locations.
Penalizing these topological errors encourages the network to learn the connectivity of neighboring voxels on the boundary and also boosts up the value of the broken locations while training the network.
The contributions of our method are following:
• We propose a topology-aware method for 3D segmentation of plant tissues using a CNN. • To train this network, we propose a new topological loss (CavityLoss) to penalize the broken connections and encourage the network to improve the segmentation performance.
• To penalize the broken connections, we propose a new method inspired by mathematical morphology to efﬁciently localize topological errors in the boundary of the prediction.
Proposed Method 2.1 Overview
State-of-the-art methods still generate discontinuities and broken connections in the 3D segmentations due to the fact that these methods do not take into account the underlying topology of the image.
In this paper, we present here a simple-yet-effective method, which integrates the geometrical information of the structures using the mathematical morphology method into the convolutional neural network.
We aim to preserve the topology and voxel connectivity on the cell boundary.
In particular, our method focuses on detecting and penalizing the topological errors in the boundary of regions.
Given an input image I and a ground truth image gt , our goal is to train the network to produce the predicted map u that has the same topology as the ground truth gt .
In other words, the predicted map u has the same Betti numbers 1 as the ground truth.
Especially in this article, for the task of 3D segmentation of plant tissues into individual cells, we focus on the number of cavities ( β 2 ) in the [Figure] image.
However, the main challenge here is that the Betti number is a discrete value, which can not be differentiable for backpropagating through the network.
Therefore, we need an alternative way to design a differentiable loss function which is able to preserve the topology information.
Instead of computing directly the number of cavities, we detect the topological errors which cause the leakage on the boundary prediction.
And then we propose a novel topological loss which targets to ﬁx topological errors of the model.
The overall framework is illustrated in Fig.
In practice, in this article, we use a 3D U-Net [3] as a network to train our topological loss because its architecture efﬁciently adapts for a 3D segmentation.
Moreover, using 3D U-Net is also adequate for us to demonstrate that
the
CavityLoss outperforms state-of-the-art methods.
Note that, theoretically our loss can be used with any other networks able to segment 3D images.
2.2 Detection of the topological errors
In this section, we propose our method to detect the topological errors.
The accuracy of the detection of the topological errors is essential for the training process to improve the performance of the segmentation, especially for the 3D bioimage structure.
We use 26-connectivity (shortly C 26 ) to represent the connectivity between neighboring voxels in 3D image.
Each voxel has 26 voxels surrounding it.
So that, 26-connected voxels are neighbors to every voxel that touches one of their faces, edges, or corners.
The identiﬁcation of the topological errors is depicted in Fig.
Note that, for better visualization, we present the process on the 2D slice of the volume, but the whole algorithm works on the 3D volume.
Our method starts with thresholding the prediction image u at 0.5.
We call u B the thresholded image.
Different from other state-of-the-art methods [6, 11] which localize the topological errors on the continuous likelihood image, our method computes the leakages on the binary image.
The main reason for that is to remove noise from the likelihood image.
Then, the task now is to train the network so that the binary image u B can be topology equivalent to the ground truth gt .
We also dilate the ground truth gt to take into account the tolerance of minor displacements in boundary location.
Indeed, minor differences in the boundary have little effect on the shape of the regions.
The dilated ground truth ˆ gt is then multiplied with the thresholded binary u B to get u M to remove the false positives in the prediction.
It is because we only focus on the leakages in the boundary region.
We then subtract the ground truth gt to prediction u M to get the different zone u Dif where the two maps gt and u M disagree.
We only choose the parts of the u Dif zone that are positive.
Note that we should not compute the loss function directly on the voxel-level disagreement because this naive metric is over-sensitive to the boundary locations.
However, we also realize that the topological errors have to belong to the different zone u Dif and can not be elsewhere.
By detecting these topological errors, we are able to preserve the topology of the structure.
To correctly extract the topological errors, we apply a segmentation method based on the watershed algorithm [12].
However, to implement the watershed algorithm, we need to initiate the set of markers.
Instead of ﬁnding local maxima of the distance map from the prediction image, we simply compute the set of connected components S from u cc , the sum of the map u M and the difference zone u Dif to get u k .
The labels S in u k also corresponds to different basin regions of the watershed algorithm [12].
We use this label image as a initiated set of markers for the watershed algorithm.
We then apply the watershed algorithm but the propagation is restricted only to the white areas of u k (which correspond to the area of u ) and we exclude the black areas as the errors can only occur in u Dif .
The detected watershed lines (zones) are the topological error.
can then T topological errors.
, we that all the of u Dif are linked with more than label in u k .
To efﬁciently implement the distance transform watershed process [12], we use Dijkstra’s algorithm to spread the distance values from the markers to other voxels in the image (voxels in the difference zone u Dif ).
We use C 26 to deﬁne the voxel connectivity in the 3D image.
We ﬁrst initialize the distance of the set of markers to zero.
We also use a priority queue to store distance values of the voxels in the propagation process.
After each iteration, we update the priority queue based on their distance costs, and update the labels of the voxels.
Different label waves meet at the boundary of the regions.
The process stops when every voxel in the image has been visited.
The propagation is also fast as it is performed in the restriction u Dif .
Finally, we are able to extract the leakages on the boundary.
2.3 Topological loss
Let’s call T the mask of the topological errors in the prediction image.
The topological loss function
L CL of the network is deﬁned as follows:
[Equation]
where (cid:12) is the element-wise product and L BCE is the binary cross entropy loss function.
The ﬁnal loss of our method, L Total , is given by:
[Equation]
where
[Equation]
L voxel is the voxel-wise loss function used to regulate the global information of the segmentation, and L CL penalizes the topological errors on the prediction image and encourages the network to recover the broken connection.
The hyper-parameter α ≥ 0 is used to balance the two losses.
Note that the mask T is spare and not ﬁxed.
Indeed, the likelihood map is updated after every iteration leading to the re-computation of the set of topological errors T .
Then, the new L CL is computed.
The training process continues until convergence.
Experiments
In this section, we test the efﬁciency of our proposed method and carry out an investigation on the impact of the hyper-parameter α on the performance of our method.
3.1 Datasets
The experiments are conducted on both ﬁxed and live publicly available microscope datasets.
These two datasets along with the ground truth were provided using different types of microscopes and differ in image quality.
• Arabidopsis thaliana ovules [15] is a ﬁxed dataset that contains 31 volumetric stacks with a voxel size of 0.075 x 0.075 x 0.235 µ m to capture all the developmental stages by using confocal laser scanning microscopy.
The dataset is split into 22 training stacks, 2 validating stacks, and a test set of 7 stacks.
• The second dataset contains 3 time-lapse videos of the development of Arabidopsis thaliana lateral root primordia (LRP) [1] using a light sheet ﬂuorescence microscope (LSFM) with a voxel size of 0.1625 x 0.1625 x 0.250 µ m .
The sizes of the movies are respectively 2048 x 1050 x 486, 1940 x 1396 x 403, and 2048 x 1195 x 566 voxels.
For this dataset, 28 stacks are used in total, where 22 stacks are used for training, 2 for validating, and 4 for testing
To improve the performance of image segmentation in the training process, we also performed an on-the-ﬂy data augmentation following the implementation in [3], including rotation, scaling, gray value augmentation, and a smooth dense deformation ﬁeld on both data and ground truth labels.
3.2 Evaluation Metrics
Boundary detection and segmentation metrics are used to evaluate the performance of our proposed method.
The evaluation metrics are presented below:
• IoU evaluates the overlap between the segmentation boundary u and the ground truth gt : IoU [Equation] gt [Equation]
• Precision and recall describe respectively the purity and completeness of the positive predictions relative to the ground truth: Precision =
TP / ( TP + FP ) and Recall = TP / ( TP + FN ) , where TP, FP, FN are respectively the true positive, false positive, and false negative.
• F-score is the harmonic mean of precision and recall : F = 2 × P ×
R/
( P + R ) .
• Adapted Rand error (ARand) is the fraction of image voxel pairs on which the two segmentations disagree: ARand = 1 −
[Equation]
, where p ij is the joint probability distribution that a randomly chosen voxel belong to region
i in the segmentation image and region j in the ground truth, (cid:80) ij p ij = 1 , s i
= (cid:80) j p ij , and t j = (cid:80) i p ij .
Greater IoU, precision, recall, and F-score scores indicate better performance.
However thinner or thicker detection does not match entirely the ground truth leading to a penalty in IoU, precision, recall, [Figure] and F-score even if the topological structure of the segmentation is preserved.
Moreover, a missing voxel which has the same penalty can even change the topology of structures.
That is the reason why we use the ARand score to capture the respect of cavity segmentation.
A smaller ARand corresponds to a method with better segmentation results.
Even if we want to get good results according to IoU, precision, recall, and F-score, our goal is to improve ARand score.
Such orders are accompanied by an upward ( ↑ ) or downward arrow ( ↓ ) in upcoming quantitative result tables.
3.3 Training Details
We pre-train the network using the linear combination of BCE and Dice and the learning rate is originally set to 2 .
10 − 4 .
Then we set the learning rate to 2 .
10 − 5 to ﬁne-tune our CavityLoss function.
The α value is chosen empirically depending on the dataset.
As such, we experimented with different values to choose the best α for the network.
Except in Sec.
3.4 the prediction is thresholded with a threshold set to 0 .
The proposed method was implemented using the Pytorch framework.
All the experiments are performed on an Nvidia Quadro RTX 8000 GPU with 48GB of VRAM.
3.4 Boundary evaluation
We conducted a comparison of the boundary detection between our proposed method with different values of the hyper-parameter α and state-of-art methods including 3D U-Net [3] and clDice [14] from the voxel-wise point of view.
To do that, we evaluate the precision/recall curves at different threshold levels of the predicted boundary probability maps.
The best performances on the LRP and
Ovule datasets are illustrated in Fig.
At a glance, our proposed method outperforms the 3D U-Net on both datasets across different values of α .
Compared to the clDice method, our method achieves better performance on the LRP dataset and comparable results on the Ovule dataset.
In conclusion, our proposed method, which is intently designed to improve the topology for image segmentation, also ameliorates the performance of boundary detection.
It is worth noting that when α increases, the boundary detection performance slightly decreases.
It is because the CavityLoss focuses on the topology which can generate some noises on the segmentation image.
3.5 Segmentation evaluation
This section focuses only on the evaluation of the segmentation of tissues into cells between our method and state-of-the-art methods.
To evaluate fairly when analyzing the strength of every method, we do not use any post-processing procedure (except thresholding the prediction at 0.5).
We then use the Adapted Rand error (ARand) for the overall segmentation quality between different segmentation methods.
The quantitative results are summarized in Tab. 1.
Our ﬁrst observation is that our proposed method achieves extremely high results for maintaining the topology of the image segmentation compared to 3D U-Net and clDice.
For example, with the value of α equal to 50, our method improves -0.3199 in ARand (in LRP dataset) and -0.093 in ARand (in Ovule dataset) compared to 3D
[Table]
[Figure]
U-Net method.
The second observation is that the clDice method improves the boundary detection compared to 3D U-Net for voxel-wise evaluations such as IoU, Precision, and F-score , but can not ﬁx the cavity problem on the 3D segmentation.
Lastly, we observe the impact of the hyper-parameter α on the segmentation results.
The segmentation scores ARand decrease (which is better) when we increase the value of α , but it also decreases slightly the voxel-wise evaluation.
This can be explained by the fact that our method focuses only on the topology of the image, in particular, the cavities in the predicted segmentation.
When the value of α is too big, our CavityLoss dominates the BCE-Dice
Loss, thereby generating some false positive noises on the segmentation results.
Fig.
4 illustrates the qualitative results of our proposed method with α equal to 50 compared to other state-of-the-art methods.
Our method which combines the voxel-wise BCE-Dice and the
CavityLoss is able to ﬁx the broken connection problem highlighted by the red rectangles.
These broken connections are because of the quality of the image ( blur and noises) and also the ﬁne-scale structure and thin connections.
Our method leverages the efﬁciency of mathematical morphology to provide topological information, thereby accurately detecting the leakages on the boundary, thus regulating the network behavior.
As a result, our method produces more accurate image segmentation and maintains the consistency of the boundary of the cell on both ﬁxed and live datasets.
Fig.
5 shows
[Figure]
Conclusion
In this paper, we propose a simple-yet-effective topology-aware method for the correct segmentation of 3D plant tissues into cells.
Our method relies on networks and especially we propose a way to train the network to detect topological issues during the learning phase, in order to focus on these errors and help the network to ﬁx them.
To do so, we introduce the CavityLoss which leverages the efﬁciency of the mathematical morphology to extract the topological information for regulating the training process of the network.
By using the well-known watershed algorithm, we are able to detect the broken connections in the boundary prediction and recover the closeness image segmentation.
For the experiments, we used the 3D U-Net but our CavityLoss can be used with any other network able to segment 3D images.
Experimenting results on the 3D ovule and lateral root datasets show that our method outperforms the state-of-the-art methods for preserving the topology of the image and also maintaining the voxel connectivities on the boundary of cells.
In the future, we will extend our method to different plant and animal organs.
Furthermore, we would design a Topology-Layer in the network to directly learn the structural representation.