I2CNet: An Intra- and Inter-Class Context Information Fusion Network for Blastocyst Segmentation

1 Introduction

Infertility is a disorder that characterized by failure to establish a clinical pregnancy following normal and unprotected intercourse within twelve months [Zegers Hochschild et al., 2009]. About 10%-15% of couples suffer from infertility around the world [Tamrakar and Bastakoti, 2019].

In-Vitro Fertilization (IVF) is one of the most common infertility treatments. During IVF process, the fertilized eggs a) Image d) Label e) U-Net f) Diff (U-Net. vs. Label)

∗Corresponding Author.

†Co-corresponding Authors.

Figure 1: The blastocyst images and labels. a) Original image; b) inter-region of ICM; c) intra-region of ICM; d) labels of different tissues for blastocyst; e) predicted map of U-Net; f) the difference map between U-Net and label.

(embryos) are first need to be cultured under constant conditions in vitro until they develop into blastocysts (about 5 days) and then high-quality embryos are selected for implantation based on morphological characteristics. However, human participation in such quality assessments may result in higher inter- and intra-observer variability. Furthermore, it is a laborious and error-prone process that requires domain knowledge and expertise. Therefore, it is necessary to explore automatic segmentation methods for different tissue regions in blastocysts.

In recent years, the development of deep learning has continuously improved the performance of medical image segmentation tasks. The most commonly used network is UNet [Ronneberger et al., 2015] and its derivative variants. There are some works that have attempted to segment blastocyst tissue, but the segmentation performance still needs to be improved. In blastocyst images, on the one hand, there are large differences in intra-class characteristics and adhesion of various tissues, as shown in the region of interest (ROI) of Fig. 1 a), b), and c); on the other hand, the inter-class structural distinction is not obvious, especially the edge area, as shown in Fig. 1 b). Since the blastocyst is a transparent object, and the image is collected under low-light conditions, this exacerbates the intra-class blur and the indistinct features at the edges. The results of traditional U-Net for blastocysts segmentation show that intra-class is no clustering, such as some areas of inner cell mass (ICM) cannot be recognized, some areas of zona pellucida (ZP) are incorrectly recognized as blastocoel, and the inter-class structures are unclear (the errors are mainly concentrated on the boundary and some hard samples' pixels), as shown in Fig. 1 a), d), e), and f).

To address the aforementioned problems, we propose an I2C module, which contains an IntrA-Class Context Module (IACCM) and an InteR-Class Context Module (IRCCM). IACCM adaptively enhances the same category and suppresses the feature weights of other categories, thereby the degree of aggregation within the class is enhanced and the segmentation results are more unified. The advantage of the IACCM is to simplify and decompose a K-category recognition task into K recognition tasks of two labels. But this also brings two problems: one is that the IACCM will be more inclined to learn easy samples, a large number of easy samples dominate the gradient. How to make up for hard samples to learn weights is particularly important, such as the pixels in the edge area account for a small proportion of the entire image and the characteristics are not sharp; the second is that although the decomposition task reduces the complexity of model learning, it lacks the structural relevance of inter-class samples. Consequently, we introduce the IRCCM, through a weighted mapping function to continuously focus on the unsolved difficulties, and these are the boundary of the category and hard samples. Meanwhile, the interaction of different classes also makes up for the shortcomings of intra-class modules.

The I2C module is equipped at each level of the network, and top-to-down continuously introduces supervision information in a coarse-to-fine manner, rendering sufficient information interaction and fusion between multi-scale features. At the high level of network, the topological structures of different categories are integrated based on semantic information. At the bottom layer of network, the image edges are processed in a more refined manner according to low-level statistical features such as color, edge, and texture. Intraclass and inter-class interactions enable the model to better learn latent features of embryos.

In summary, the main contributions of this work can be summarized as follows:

• A general architecture framework named I2CNet is proposed in this work, which reveals how to leverage information of intra-class and inter-class to consistently boost the segmentation performance of blastocyst image.
• We design a simple and effective IACCM to enhance the aggregation of the same class and a novel IRCCM to capture the relation information of different class, respectively. Experimental results demonstrate the effectiveness of our method. Especially, the IRCCM contains a weighted mapping function, which can promote the represents of boundary and hard pixels.
• The top-to-down feature enhancement pathway which couples the backbone encoder increases the representation power of feature maps with intra-class enhancement

and inter-class enhancement in a coarse-to-fine manner.

2 RelatedWork

In recent years, there are many works that have attempted to segment blastocyst image [Filho et al., 2010; Khan et al., 2016], but mainly focus on the identification of ICM and trophectoderm (TE).

The semi-automatic segmentation method includes the level set idea [Filho et al., 2012]. By defining the initial contour, the gradient information is applied to segment the ICM and TE. However, in clinical practice, the zona pellucida of the blastocyst is ruptured and the ICMwill be outside the zona pellucida. In this method, it is difficult to define contour information and the recognition effect is not significant. Moreover, it requires human intervention. On this basis, Singh et al. [Singh et al., 2014] applied level set combined with the correction of Retinex to automatically identify blastocysts, mainly for TE segmentation. It achieved an accuracy of 87.8% and a recall rate of 78.7%. Rad et al. [Rad et al., 2017] utilized image texture features (Gabor and DCT features) and combined level sets to determine ICM boundaries, obtaining a Jaccard index of 70.3%. Saeedi et al. [Saeedi et al., 2017] also applied texture features, but combined k-means clustering and watershed to identify ICM and TE, where ICM and TE achieved Jaccard index of 71.1% and 63%, respectively. Kheradmand et al. [Kheradmand et al., 2016] combined Discrete Cosine Transform (DCT) and used a two-layers neural network for feedback propagation to identify ZP, TE and ICM, but only obtained Jaccard indices of 47.7%, 58.9% and 67.4%. Kheradmand et al. [Kheradmand et al., 2017] used a Fully Convolutional Neural (FCN) network to segment the ICM and achieved a Jaccard Index of 76.5%. The ICM recognition methods that FCN-based [Leahy et al., 2020] and U-Net based [Rad et al., 2018] outperform [Kheradmand et al., 2016; Rad et al., 2017; Saeedi et al., 2017; Kheradmand et al., 2017]. Yu and Koltun [Yu and Koltun, 2015], Rad et al. [Rad et al., 2019] proposed a U-Net variant by adding dilated convolution, which only increases the perception field of a single pixel, but not pay attention to the attribute relationship between pixels, and Jaccard index is below 82%.

The above methods are either dedicated to the segmentation of a single tissue or identification of several tissues. It is necessary to focus on the overall identification of a single tissue and the structural associations of all tissues. There is still a research gap of how to improve the segmentation performance of blastocyst.

3 Method

An overview of the proposed framework is illustrated in Fig. 2. The I2CNet incorporates the intra-class contextual information and inter-class contextual information for blastocyst image segmentation.

3.1 Formulation of I2CNet

Given a set of images S, Let S = {(Xm, Ym),m = 1, · · · ,M}, where Xm denotes the original input embryo image and Ym denotes the corresponding ground truth. Since each image Xm is processed separately, the subscript m is omitted for simplicity in the following sections.

We define the encoding block as E = {En(·), n = 1, 2, · · · ,N}, where N is the numbers of encoding blocks. En(·) contains two convolutional layers, which contains two Batch Normalization (BN) layers, two ReLU activation layers and a max pooling layer. We first use E to project the pixels in X into a non-linear embedding space so that we can obtain the pixel representations RnE

where RnE is a matrix of size C × H 2n × W 2n , C is the channel of RnE .

We use F(·) squeeze RnE to representations RnF that with K channel (K is the number of class), where F(·) contains two 1×1 convolutional layers. The process can be formulated as follows

where the Rn I2C is the features representation that aggregate the intra- and inter-class information from the whole image in different scale by I2C module. I2C module F I2C n-1 (·) contains intra-class context model F intra n-1 (·) and inter-class context model F inter n-1 (·), as shown in Fig. 3.

After squeezing the channel of representations by F(·), we use L(·) to get the supervised logit map RnL of K × H 2n × W 2n scale images

where L(·) contains a convolutional layer, BN and ReLU. Following above expectation, as indicated in Fig. 2, before gettinng the aggregate information of the intra- and intermodule, we get the probability map from the lower level logit map RnL.

To match the scale of high level representations, we use upsampled function to get the probability map P n with size of (K ×H′ ×W ′ ,H′ = H 2n-1 ,W ′ = W 2n-1 )

where σ(·) is the sigmoid function. Then, I2C module F I2C n-1 (·) can use the P n to enhance the Rn-1 E features (K × H′ ×W ′ ,H′ = H 2n-1 ,W ′ = W 2n-1 ):

Finally, the output predicted map in different scale is define as On:

3.2 Intra-Class Context Module

In I2C module, intra-class context module Fintra is designed to capture the context information from the same class of the whole image, as shown in the top of Fig. 3. There are feature enhancements both in the same scale and across different scales. Before intra-scale enhancement, every level logit map passes a supervised module before fusion. Then, these logit maps are activated with neighboring high level features to focus on the context information from the same class.

To make the adjacent high-level features match the probability map of the K-class tissues, we orderly grouped Rn-1 E into K groups, (Rn-1 E )k is a matrix of size Ck× H 2n-1 × W 2n-1 , where Ck denotes the number of channels belonging to the category k. In this way, the number of parameters can be effectively reduced

For the convenience of presentation, we define the (P n )k is the logit map belonging to category k. To aggregate contextual information from the same class, we compute the region representation for each class k as follows:

Eventually, we orderly concatenate the all (Rn-1 intra )ki and use a 1 × 1 convolution to get the final Rn-1 intra features. The Rn-1 intra features contain the context information from the same class of the whole image. By grouping, we decompose a multi-category recognition task into multiple two-category recognition tasks, thereby reducing the learning difficulty of the model, allowing the model to focus on learning features of the same category without interference from other categories. Moreover, the amount of parameters is reduced to a certain extent.

3.3 Inter-Class Context Module

Although IACCM effectively solves the problem of feature representation and learning within the same class, it ignores the inter-class interconnection. As shown in Fig. 1 f), it can be seen that the predictive error between classes accounts for a relatively large proportion, including some pixels with unobvious structural features. To address this problem, interclass context module Finter is designed to deal with the relation information of different class and some hard pixels , as shown in the bottom of Fig. 3.

The inter-class and intra-class modules share the same high-level features and follow the same grouping principle. The intra-class grouping is equal to the number of tissues categories, but the inter-class grouping considers the correlation between tissues, which is K -1 group. And the tissues order of the blastocyst from inside to outside is ICM, Blastocoel, TE, ZP, and Background. So we only considered four organizational relationships: 1, 1+2, 1+2+3, and 1+2+3+4, where 1, 2, 3, 4 represent ICM, Blastocoel, TE, ZP, respectively.

To aggregate the context information from the inter class, we calculate the region representation as follows:

where U(·) is a weighted mapping function, which can enhance the represents of the border and hard pixels. According to Eq. 11, we can get the structural feature map of the correlation between the inside-out (ICM to ZP) tissues.

In IRCCM, according to the characteristics of cross entropy loss, as the probability approaches 0.5, the greater the loss, which means that the sample is harder to recognize. For embryo samples, the samples of the boundary between different tissues is small. Moreover, the blastocyst is a transparent object and the image is collected under low-light conditions, resulting in unclear boundary features. Therefore, we designed a weighted mapping function U(·) to redistribute the probabilities weights of each category, that is, the weight value of the sample pixels that are not easy to identify is close to 1+α(1-e-( 12 )β ), and the weight value of the sample pixels that are easy to identify is close to 1, which makes the model more effective for learning and reasoning at edges and hard pixels. The U(·) is defined as:

where α and β are the activation factors, and set as 1 and 2 in our experiment.

Finally we orderly concatenate the (Rn-1 inter )ki and use convolution with kernel size 1×1 to get the final Rn-1 inter features. The Rn-1 inter features contain the context information from the different class of the whole image.

3.4 Loss Function

In this work, we use cross entropy loss to supervise logit map, it can be defined as Eq. 13:

where Y n S represents the one-hot label (K×H×W) scale to the size of K × H 2n × W 2n , and (RSn L)K× H 2n × W 2n is calculated as Eq. 14:

Lce denotes the cross entropy loss and ∑i,j (·) denotes that summation is calculated over all locations on the input image X. The total cross entropy loss define as:

4 Experiments

4.1 Experimental Setup

Dataset and Ground Truth. The dataset is introduced in [Saeedi et al., 2017], which is the only public dataset on human blastocyst. This benchmark consists of 249 blastocyst images and their Ground Truth (GT) masks provided by the Pacific Centre for Reproductive Medicine (PCRM). The training set comprises 199 images (80%) and the testing set contains 50 images (20%).

Training Details. In this work, the number of encoding blocks is 5. The backbone and two integrated context modules are initialized randomly by kaiming. The model was trained with 2000 epochs and the Adam was used as the optimizer. Initial learning rate and weight decay were set to 1 × 10-4 and 1 × 10-5 . Poly learning rate policy with factor (1 - epoch total epoch-epoch )0.9 is performed for training. Synchronized batch normalization implemented by PyTorch is enabled during training. For data augmentation, we process each sample with random scaling in the range of [0.5, 2.0], random rotating in the range of [-15, 15] degrees, random cropping and left-right flipping in training stage. The input image size is 256×256. By default, no tricks (e.g. multi-scale and multi-rotation with flipping testing) are adopted during the testing stage.

Reproducibility. The proposed framework is implemented on PyTorch (version ≥ 1.3) and trained on four NVIDIA Tesla V100 GPUs with a 32 GB memory per-card. And all the testing procedures are performed on a single NVIDIA Tesla V100 GPU. To provide full details of our framework, our code will be made publicly available.

Evaluation Metrics. In order to evaluate the performance of the network, there are five commonly used indicators: Accuracy, Precision, Recall, Dice Coefficient and Jaccard Index [Csurka et al., 2004; Zhu et al., 2016; Thoma, 2016; Taha and Hanbury, 2015].

4.2 Ablation Study

Effectiveness of IACCM. In embryo images, there is a phenomenon of intra-class continuous aggregation, in which the inner cell mass is the most core continuous region, while other tissues are band-like regions but also continuous. It is very important to ensure that the model can learn to understand and recognize the features of the same category more accurately during the training process. As shown in the third column of Fig. 4, we can see that the proposed IACCM can effectively focus on learning the features of the same category and model the continuous dependencies of pixels in the same category, making the intra-class tissues more clustered. In Table 1, we can see that the addition of IACCM improves the overall performance of the network, in which the Dice and Jaccard indicators increase by 2.73% and 4.4% respectively. The experimental results show that the IACCM is effectiveness for increasing intra-class segmentation performance.

Effectiveness of IRCCM. In medical images, especially for embryo images, focusing on the logical relationship between different tissues is also very helpful for accurate segmentation of embryo tissues. Since the tissue types of embryo cells from inside to outside are arranged according to developmental law, such as the inner circle of the trophoblast must be connected to the cyst and the outer circle must be connected to the zona pellucida. Therefore, we designed an IRCCM to make the model learn and express the features of embryo images more efficiently by modeling inter-class relationships progressively. In particular, we add a weighted mapping function U(·) to increase the weights of boundaries and hard samples. From Table 1, we observe that U(·) improves the performance of IRCCM, increasing the Accuracy from 93.04% to 93.58%, the Dice from 90.19% to 91.30% and the Jaccard from 83.28% to 84.77%. The introduction of IRCCM with U(·) improves the backbone performance, which improves Dice by 2.77% and Jaccard by 4.49%. The visualization results are shown in the fourth column of Fig. 4, clearly showing that the IRCCM is effectiveness for improving the inter-class recognition error. These improvements suggest that introducing IRCCM component can model the relationship between classes well and increase the learning ability of the model.

Effectiveness of IACCM& IRCCM. In this work, we integrate IACCM and IRCCM to get the I2C module. As illustrated in Table 1, we can find that combining IACCM and IRCCM outperforms the backbone by 4.97% in terms of Jaccard. The improvement is much higher than applying single IACCM (4.97% v.s.4.40%) or single IRCCM (4.97% v.s. 4.49%). Moreover, from the fifth column of Fig. 4, we can see that after combining the IACCM and IRCCM modules, not only the intra-class is more aggregated, but the interclass relationship is also more reasonable and the boundary is clearer. These results indicate that IACCM and IRCCM can complement and promote each other, which well demonstrates the reliability of the basic motivation, and the effectiveness of the designed framework in this paper.

4.3 Comparison with State-of-the-Art Methods

Table 2 lists the blastocyst segmentation performance comparison between the proposed framework and related methods in the literature. Among the compared methods, there FLOPs(G)

Figure 6: Comparison of Parameters, FLOPs and Jaccard for different methods. Note that parameters and FLOPs are calculated on a 256 ×256 input image.

are models including the classic CNNs and the more popular transformers. Deeplabv3+, SETR, and Pranet apply pretrained models, which are noted in the back brackets. Others are trained from scratch and we have optimized their hyperparameters. It can be seen from the Table 2 that the proposed I2CNet achieves the state-of-the-art performance compared with other methods, among which Accuracy, Precision, Recall, Dice and Jaccard achieve 94.14%, 91.18%, 92.32%, 91.56%, 85.25%, respectively. Fig. 5 shows some qualitative results for U-Net, Blast-Net and we proposed I2CNet. As seen, the results show that our proposed I2CNet achieves the most similar results to the ground truth. These experimental results demonstrate that the effectiveness of the proposed method.

In addition, Fig. 6 shows the parameters, FLOPs and Jaccard for the comparison methods and our proposed method. The size of the circle represents the parameter amount of the model (the larger circle means a larger number of parameters in the model). It can be intuitively seen that I2CNet has achieved the best Jaccard index, meanwhile the parameters and FLOPs are fewer, which further illustrates the superiority of our proposed method.

5 Conclusions

In this work, we studied the intra-class and inter-class relationship of embryo image, and proposed an IntrA-Class Context Module and an InteR-Class Context Module. The IACCM transforms a multi-class recognition task into multiple binary classification tasks, and employe the supervised probability map to guide the learning of high-level features, which not only reduces the complexity of the model but also improves the continuity of intra-class tissues. The IRCCM transmits the relationship probability map of different tissues to the high-level to guide the learning of the model. Meanwhile, the probability map is remapped and assigned new weights to enhance the weights of boundary and hard samples, so that the model can better focus on valuable sample pixels instead of blindly learning simple samples. The proposed I2CNet integrates both IACCM and IRCCM to achieve state-of-the-art performance. Finally, we qualitatively and quantitatively verify the effectiveness of our proposed model through ablation experiments and comparative experiments with other state-of-the-art methods.

Acknowledgments

This work was supported in part by the Beijing Natural Science Foundation under Grant Z200024; in part by Hefei Innovation Research Institute, Beihang University; and in part by the University Synergy Innovation Program of Anhui Province under Grant GXXT-2019-044.