Learning interactive multi‐object segmentation through appearance embedding and spatial attention

IET Image ProcessingVolume 16, Issue 10 p. 2722-2737 ORIGINAL RESEARCHOpen Access Learning interactive multi-object segmentation through appearance embedding and spatial attention Yan Gui, Corresponding Author Yan Gui [email protected] orcid.org/0000-0001-8323-4571 School of Computer and Communication Engineering, Changsha University of Science and Technology, Changsha, China Hunan Provincial Key Laboratory of Intelligent Processing of Big Data on Transportation, Changsha University of Science and Technology, Changsha, China Correspondence Yan Gui, School of Computer and Communication Engineering, Changsha University of Science and Technology, Changsha, China. Email: [email protected]Search for more papers by this authorBingqiang Zhou, Bingqiang Zhou School of Computer and Communication Engineering, Changsha University of Science and Technology, Changsha, China Hunan Provincial Key Laboratory of Intelligent Processing of Big Data on Transportation, Changsha University of Science and Technology, Changsha, ChinaSearch for more papers by this authorJianming Zhang, Jianming Zhang [email protected] orcid.org/0000-0002-4278-0805 School of Computer and Communication Engineering, Changsha University of Science and Technology, Changsha, China Hunan Provincial Key Laboratory of Intelligent Processing of Big Data on Transportation, Changsha University of Science and Technology, Changsha, ChinaSearch for more papers by this authorCheng Sun, Cheng Sun School of Mathematics and Statistics, Hunan Normal University, Changsha, ChinaSearch for more papers by this authorLingyun Xiang, Lingyun Xiang [email protected] orcid.org/0000-0001-7396-0908 School of Computer and Communication Engineering, Changsha University of Science and Technology, Changsha, China Hunan Provincial Key Laboratory of Intelligent Processing of Big Data on Transportation, Changsha University of Science and Technology, Changsha, ChinaSearch for more papers by this authorJin Zhang, Jin Zhang [email protected] orcid.org/0000-0002-7464-2247 School of Computer and Communication Engineering, Changsha University of Science and Technology, Changsha, China Hunan Provincial Key Laboratory of Intelligent Processing of Big Data on Transportation, Changsha University of Science and Technology, Changsha, ChinaSearch for more papers by this author Yan Gui, Corresponding Author Yan Gui [email protected] orcid.org/0000-0001-8323-4571 School of Computer and Communication Engineering, Changsha University of Science and Technology, Changsha, China Hunan Provincial Key Laboratory of Intelligent Processing of Big Data on Transportation, Changsha University of Science and Technology, Changsha, China Correspondence Yan Gui, School of Computer and Communication Engineering, Changsha University of Science and Technology, Changsha, China. Email: [email protected]Search for more papers by this authorBingqiang Zhou, Bingqiang Zhou School of Computer and Communication Engineering, Changsha University of Science and Technology, Changsha, China Hunan Provincial Key Laboratory of Intelligent Processing of Big Data on Transportation, Changsha University of Science and Technology, Changsha, ChinaSearch for more papers by this authorJianming Zhang, Jianming Zhang [email protected] orcid.org/0000-0002-4278-0805 School of Computer and Communication Engineering, Changsha University of Science and Technology, Changsha, China Hunan Provincial Key Laboratory of Intelligent Processing of Big Data on Transportation, Changsha University of Science and Technology, Changsha, ChinaSearch for more papers by this authorCheng Sun, Cheng Sun School of Mathematics and Statistics, Hunan Normal University, Changsha, ChinaSearch for more papers by this authorLingyun Xiang, Lingyun Xiang [email protected] orcid.org/0000-0001-7396-0908 School of Computer and Communication Engineering, Changsha University of Science and Technology, Changsha, China Hunan Provincial Key Laboratory of Intelligent Processing of Big Data on Transportation, Changsha University of Science and Technology, Changsha, ChinaSearch for more papers by this authorJin Zhang, Jin Zhang [email protected] orcid.org/0000-0002-7464-2247 School of Computer and Communication Engineering, Changsha University of Science and Technology, Changsha, China Hunan Provincial Key Laboratory of Intelligent Processing of Big Data on Transportation, Changsha University of Science and Technology, Changsha, ChinaSearch for more papers by this author First published: 01 May 2022 https://doi.org/10.1049/ipr2.12520AboutSectionsPDF ToolsRequest permissionExport citationAdd to favoritesTrack citation ShareShare Give accessShare full text accessShare full-text accessPlease review our Terms and Conditions of Use and check box below to share full-text version of article.I have read and accept the Wiley Online Library Terms and Conditions of UseShareable LinkUse the link below to share a full-text version of this article with your friends and colleagues. Learn more.Copy URL Share a linkShare onFacebookTwitterLinkedInRedditWechat Abstract Deep learning approaches to interactive image segmentation are typically formulated as a binary labeling problem. A model trained to make predictions within a fixed set of labels (i.e., foreground and background labels) cannot be used to directly predict the binary masks of multiple objects of interest, which greatly limits its flexibility and adaptivity. The use of different classes of clicks as input is opted for and the first end-to-end learning model for multi-object segmentation, based on a new designed neural network, is developed. The network consists of a visual feature extractor, a recurrent attention module and a dynamic segmentation head, extracts user click-adapted appearance embedding features and spatial attention features, and then learns to transform this information into a segmentation of multiple objects. It is also proposed to train the network using a joint loss function, taking the embedding learning into account for segmentation. Comprehensive experiments are conducted on three benchmark datasets to demonstrate the effectiveness of the proposed method. It performs favorably against state-of-the-art approaches on the multiple object segmentation task, for example, with 0.15 s per image, 0.06 s per object and mean IoU & F1 score of 84.90% on Pascal VOC 2012 validation set. It is further shown that the method can be used in numerous vision applications such as image recoloring and colorization. 1 INTRODUCTION Interactive image segmentation has proven to be a vital technology and influenced various computer vision and computer graphics applications, including image/video editing, scene understanding, medical image analysis, and object retrieval. Interactive image segmentation is to take a small amount of user interactions through a set of clicks, scribbles, or bounding boxes and segment object(s) of interest in an image from its background. It turns out to be a very challenging task especially in the situations of significantly overlapping foreground and background appearances, weak edges, intricate texture, cluttered background, poor illumination, and so on. Most traditional interactive segmentation methods [1-8] are generally formulated as a graph cut optimization problem and solved by minimizing a Markov Random Field energy through a min-cut/max-flow algorithm [9]. These methods typically adopt low-level features (color, texture etc.) as main cues to infer segmentation. Recently, various methods [10-15] approach the problem with deep learning architectures such as convolutional neural networks (CNNs), thanks to the availability of large-scale datasets such as Pascal VOC 2012 [16] and MS-COCO 2014 [17] that enable the training of CNN-based segmentation models. By working with high-level features produced by CNNs, these deep-learning based methods have much better understanding of objectness and semantics, which brings a great improvement in segmentation accuracy and significantly reduces the amount of user interaction required to generate quality results. However, all the above methods mainly process segmentation of individual objects. When applying these methods to images with multiple objects, the foreground prediction needs to be performed once for each object, which is time consuming and tedious for users. A variety of directions have been explored, such as segmenting the repeated scene elements using graph cuts techniques [8], annotating all objects and background regions in an image [18, 19], or extending the detected objects from instance-level to pixel-level [20]. Although various successful examples have been shown, these methods have not reached a point that can be expected to simultaneously cut out multiple objects of interest based on user inputs. In this paper, we devise a unified, end-to-end trainable neural network architecture that tackles the multi-object segmentation problem in an interactive framework. As shown in Figure 1, our network starts with visual feature extractor (VFE). It extracts deep visual features from the image and applies an embedding head to compute appearance embeddings. We then develop a recurrent attention module (RAM) to handle a sequence of user interactions. The RAM learns competitions between and within user interactions and produces an attention map for each object and the background, which encodes the location and appearance of objects previously found in the image. This enables our model to understand the users' intentions well. The final component of our network is a dynamic segmentation head (DSH). It fuses the features from the outputs of the VFE and RAM by a feature enhancement mechanism, and uses them to predict the multiple objects of interest. For network training, we use a joint loss on many of (image, user interactions) pairs to jointly train for embedding learning and multi-object segmentation, where the user interaction data is generated by combining adaptive object scale sampling strategies. Such joint training can yield both better performance and faster convergence. We provide extensive analysis and evaluation on the influence of each component of the network through the ablation studies. We also conduct comprehensive experiments on challenging benchmark datasets to demonstrate the effectiveness of our proposed method. FIGURE 1Open in figure viewerPowerPoint Overall architecture of Interactive Multi-Object Segmentation Network, which consists of a visual feature extractor, a recurrent attention module and a dynamic segmentation head. Given an input image and a sequence of user interactions, our algorithm first transforms click sets (denoted as color dots) into distance maps and directly maps the image and the distance maps into the binary segmentation masks, corresponding to the annotated objects and the background Our main contributions are summarized as follows: 1) We are the first to directly solve the multi-object segmentation problem by combining user interactions with the current deep learning models. This is a generalization of interactive image segmentation. 2) We adopt the VFE and RAM which allow our network to robustly acquire the prior knowledge about objects appearance and fully exploit the valuable information in user interactions to guide the segmentation process towards the desired output. 3) We achieve a new state-of-the-art performance for multi-object segmentation on the Pascal VOC 2012 validation set, with a mean IoU and F1 score of 84.90% when training on Pascal VOC 2012 and an average 0.15 s per image and 0.06 per object when running on RTX 2080 Ti GPU. The remainder of the paper is organized as follows. In Section 2, we give an extensive overview of the related work. In Section 3, we detail the proposed network for multi-object segmentation. Section 4 introduces the random sampling strategies to generate user interaction data. Experimental results and comparisons are presented in Section 5, and Section 6 concludes the paper. 2 RELATED WORKS Interactive object segmentation There are many methods proposed for interactive image segmentation, such as graph cut [1], random walks [2], Grabcut [3], geodesics [4], the combination of graph cut and geodesics [5], and methods with more powerful random field models [6]. However, all the above methods are just suitable for single object segmentation. Cheng et al. proposed RepFinder [7] which explores geometry regularity and detects the repeated elements in the image. The RepSnapping system of Huang et al. [8] used a new optimization model to quickly segment the repeated elements. However, these methods [7, 8] rely on appearance similarity between repeated elements, thereby limiting their use for the multiple object segmentation task. With the recent development of deep learning techniques, several researchers have applied convolutional neural networks (CNNs) [21] to gain features with higher-level understanding for image and user interactions to improve interactive image segmentation. Xu et al. [10] proposed iFCN. This method transforms positive and negative clicks into distance maps. A fully convolutional network (FCN) [22] was trained to predict the foreground and background from the concatenated image and distance maps. Subsequent works build upon these findings and go further by introducing newer CNN architectures [11-15], iterative training procedure [23, 24], user inputs encoding [25], polygon vertex predictions [26] and backpropagating refinement scheme [27]. All these methods have a commonality that they input the RGB image plus two extra channels for foreground object and background, and output a binary mask of the selected object. Very recently, most relevant to our work are methods that address the task of interactive full image segmentation. Andriluka et al. [18] introduced Fluid Annotation, which uses Mask R-CNN to generate a proposal set of segments and then provides an efficient interface for the annotator to obtain accurate segmentations. Agustsson et al. [19] used Mask R-CNN [20] and extreme points [28] to extract RoI features, making the segmentation much faster and more memory efficient. These methods iterate between the annotator providing corrections and the machine model refining the segmentation. In contrast, our method addresses segmenting multiple objects by combining a segmentation model and user input within an interactive framework. Furthermore, our method focuses on handling multi-object segmentation, rather than blindly producing full annotations that might be redundant for certain applications, such as image and video editing, image retrieval [29]. Object instance segmentation Our work is also related to object instance segmentation which aims to produce a binary segmentation mask for each object instance in an image. The recent top-performing methods for object instance segmentation often utilize a powerful detector to generate high-quality bounding boxes and then predict pixels within each box. Dai et al. [30] used a cascade of networks to merge bounding boxes, segmentation masks and category information. He et al. [20] proposed Mask R-CNN which extends Faster R-CNN [31] by adding a parallel segmentation branch. Lee et al. [32] proposed an anchor-free one stage object detector that directly predicts boxes. Cao et al. [33] improved object detection by applying dense local regression for precise object localization. Zhang et al. [34] provided the lost details with fine-grained features during the instance-wise segmenting process in order to refine instance masks. However, these methods can fail if there is more than one object instance inside of the box. In this work, we manually annotate several objects in the image for initialization, without the bounding-box detection step. Other recent works are based on deep metric learning [35] and train a network to output embeddings that are subsequently clustered into individual instances with a post-processing operation. Schroff et al. [36] proposed to embed a face image into an embedding vector by a network trained with a triplet loss. Brabandere et al. [37] utilized a discriminative loss to optimize distances between individual pixels in an image. Fathi et al. [38] proposed learning an instance-sensitive embedding for each pixel. Ying et al. [39] applied embedding coupling for the task of instance segmentation. These methods use the learned embeddings directly to produce the instance masks, however the accuracy of the embeddings often affects the final segmentation quality. We learn the embeddings and use them to enhance the features of the convolutional network (see Section 3.2), which are useful for the final segmentation. Another line of research is to employ recurrent neural networks to perform object detection and segmentation. For example, Stewart and Andriluka [40] extracted features using a CNN and then used a LSTM [41] to detect different human faces one at a time. Romera-Paredes and Torr [42] used convolutional LSTM (ConvLSTM) [43] to produce the object instances sequentially. Salvador et al. [44] proposed an encoder-decoder architecture based on RNN that predicts a mask for each object instance of the image at each iteration. In our work, we use a recurrent module based on ConvLSTM to deal with variable number of objects and memorize relevant states of objects. Additionally, Wang et al. [45] adopted ConvLSTM to produce an attention map that highlights the visually important regions, implicitly providing the prior knowledge of the objects to be segmented. Inspired by this, we use the attention maps, which are extracted from the given user interactions by a recurrent module, as a guidance information to a segmentation head (see Section 3.2). This allows our network to fully exploit the information in user interactions during the procedure of the multi-object segmentation. 3 PROPOSED METHOD Overview In this work, we extend the deep-learning based interactive object segmentation [10] to multi-object segmentation. Our proposed method is to use a trained model to segment multiple objects of interest in an image from its surrounding background, treating user inputs as ground-truth labels, which means the segmentation matches the objects intended by the user. More specifically, given an image I ∈ R w × h × 3 $I \in {\mathbb {R}^{w\times h\times 3}}$ and a sequence of user interactions { s i } i = 1 N $\lbrace s_i\rbrace _{i=1}^N$ , where s i $s_i$ is a set of clicks that are made on the ith object and N is the number of objects in the image, we first employ the Euclidean distance transformation in [10] to transform each click set s i $s_i$ to an individual distance map d i ∈ R w × h $d_i\in {\mathbb {R}^{w\times h}}$ . Then, our proposed method learns an interactive multi-object segmentation network model as a mapping function F : R w × h × ( 3 + N ) ↦ { 0 , 1 } w × h × ( N + 1 ) $F:{\mathbb {R}^{w\times h \times (3+N)}}\mapsto \lbrace 0,1\rbrace ^{w\times h\times (N+1)}$ that directly maps the image I and the distance maps { d i } i = 1 N $\lbrace d_i\rbrace _{i=1}^N$ into the binary segmentation masks { M i } i = 1 N + 1 $\lbrace M_i\rbrace _{i=1}^{N+1}$ , corresponding to the annotated objects and the background: { M i } i = 1 N + 1 = F I , { d i } i = 1 N . \begin{align} \lbrace M_i\rbrace _{i=1}^{N+1} = F\left(I,\lbrace d_i\rbrace _{i=1}^N\right). \end{align} (1) To learn the end-to-end mapping, the model is trained on a dataset with instance-level object annotations. As several valid segmentations can be produced for one image, the ground truth instance masks are different when different user interactions are provided. Once the model is trained, the parameters are fixed and require no optimization when tested on a new dataset. Every time the user clicks are added or removed, the segmentation is updated. Next, we first describe in detail our designed network architecture for interactive multi-object segmentation, including visual feature extractor (VFE), recurrent attention module (RAM) and dynamic segmentation head (DSH). Then, we give the network training process. Interactive multi-object segmentation network 3.2.1 Visual feature extractor We use the VFE to map each image pixel to an n-dimensional embedding vector, in which the pixels with similar appearance have organized themselves in high-dimensional embedding space, in order to associate pixels with the correct objects. The VFE includes two parts: a backbone network and an embedding head, as shown in Figure 2. We take ResNet101 based DeepLab-v3 as the backbone, where the last two blocks of ResNet101 take stride as 1 and the final classification layer is removed, to extract deep visual features with the same size of the input image. In more detail, given the input image I with the size of w × h $w\times h$ , the output f = D e e p L a b v 3 ( I ) $f=DeepLabv3(I)$ of the backbone network is a w × h × 256 $w \times h \times 256$ feature representation. FIGURE 2Open in figure viewerPowerPoint Architecture of the visual feature extractor, which consists of a backbone network and an embedding head. The module extracts 128-dimensional embedding vectors from the image. Visualization of the embedding vectors by randomly projecting the 128-dimensional vectors into RGB space For the embedding head (EH) shown in Figure 2, it consists of a stake of two depthwise separable convolution layers [48] with kernel size of 3 × 3 for the depthwise convolutions and widths of 256 and 128. The embedding head is applied on the top of the deep visual features obtained from the backbone network, and produces appearance embeddings e = E H ( f ) $e=EH(f)$ , with the size of w × h × 128 $w \times h \times 128$ . As a result, each image pixel p is represented as 128-dimensional feature vector e p $e_p$ . The output embeddings are learned such that pixels of the same object in the image are close to each other in the embedding space and those of different objects are far apart. To this end, we train the VFE by minimizing an embedding loss based on distance metric learning principles [35], which will be detailed in Section 3.3. 3.2.2 Recurrent attention module The RAM in our network aims to learn the competitions between and within user interactions and compute attention maps for the objects and background, in order to gain a better understanding of the locations of the objects (foreground). As shown in Figure 3, it is composed of a memory unit and a few depthwise separable convolution layers. The memory unit is a two-layer recurrent architecture of ConvLSTMs, where each ConvLSTM layer has a width of 64. For object i, the RAM uses the distance map d i $d_i$ transformed from s i $s_i$ , the deep visual features f and the hidden states h 1 , i − 1 $h_{1,i-1}$ , h 2 , i − 1 $h_{2,i-1}$ from the previous object i − 1 $i-1$ to compute the new hidden states h 1 , i $h_{1,i}$ , h 2 , i $h_{2,i}$ , which store the location and appearance information about the previously found objects in its hidden states, leading to a competing mechanism. This can be specified by the following update equations: h 1 , i = ConvLSTM 1 ( [ d i , f ] , h 1 , i − 1 ) , \begin{align} h_{1,i} & = \mathrm{ConvLSTM_1}([d_i, f],h_{1,i-1}),\end{align} (2) h 2 , i = ConvLSTM 2 ( h 1 , i , h 2 , i − 1 ) , \begin{align} h_{2,i} & = \mathrm{ConvLSTM_2}(h_{1,i}, h_{2,i-1}), \end{align} (3)where C o n v L S T M k ( · ) $ConvLSTM_k(\cdot )$ denotes the kth ConvLSTM layer. [.,.] represents the concatenation operation. The new hidden states h 1 , i $h_{1,i}$ , h 2 , i $h_{2,i}$ are computed at full resolution. For the first object (i.e., i = 1 $i=1$ ), the hidden states h1, 0, h2, 0 of the two-layer ConvLSTM structure are all zero matrixes, which represent there are no previous hidden states for this object. FIGURE 3Open in figure viewerPowerPoint The unfolding structure of the recurrent attention module, which takes the distance maps as input and outputs attention maps of objects and the background On the top of this, we add two depthwise separable convolution layers with a kernel size of 3 × 3 and widths of 64 and 1, which map the width of the hidden state to 1. Then, we apply a per-pixel sigmoid. Hence, the resultant pixel values are all in the interval [0, 1], as required. We extract an attention map A i ∈ [ 0 , 1 ] w × h × 1 $A_i \in [0,1]^{w \times h \times 1}$ , which represents the probability of a particular location belonging to the object. This is expressed by: A i = f s i g m o i d ( DSConvs [ 64 , 1 ] ( h 2 , i ) ) , \begin{align} A_i = f_{sigmoid}(\mathrm{DSConvs_{[64,1]}}(h_{2,i})), \end{align} (4)where h 2 , i $h_{2,i}$ is the hidden state of the last ConvLSTM layer. f s i g m o i d ( · ) $f_{sigmoid}(\cdot )$ is the sigmoid function. D S C o n v s [ 64 , 1 ] $DSConvs_{[64,1]}$ represents the depthwise separable convolution, where the value of each element in the subscript [64,1] indicates the width of the corresponding depthwise separable convolution layer. Hence, given the distance maps { d i } i = 1 N $\lbrace d_i\rbrace _{i=1}^N$ transformed from the user interactions { s i } i = 1 N $\lbrace s_i\rbrace _{i=1}^N$ , the RAM is expanded N iterations with shared weights, as the unfolding structure shown in the Figure 3, and outputs N phased attention maps { A i } i = 1 N $\lbrace A_i\rbrace _{i=1}^N$ . Considering that the distance maps are too simple and do not utilize any image information, we concatenate the distance map ( d i $d_i$ ) of each object with the deep visual features (f) obtained from the backbone network of the VFE at each iteration, as shown in the left part of Figure 3, in order to contain high-level visual information exploited for object localization. In this way, the attention maps output by the RAM can retain the position and object details, as shown in the right part of Figure 3, providing a richer representation of user interactions. Additionally, we need to compute the attention map for the background, although we do not provide user clicks to indicate the background. This is because there exists a competition between all annotated objects and background. For this purpose, we also use three depthwise separable convolution layers with a kernel size of 3 × 3 and widths of 64, 64 and 1, as shown in the middle part of Figure 3. The extra branch takes as input the hidden state h 2 , N $h_{2,N}$ of the last ConvLSTM layer of the Nth iteration, which encodes the location and appearance information of all annotated objects in the image, and feeds it into a sigmoid transformation. The attention map of the background is computed as follows: A N + 1 = f s i g m o i d ( DSConvs [ 64 , 64 , 1 ] ( h 2 , N ) ) . \begin{align} A_{N+1} = f_{sigmoid}(\mathrm{DSConvs_{[64,64,1]}}(h_{2,N})). \end{align} (5)Finally, the attention maps { A i } i = 1 N + 1 $\lbrace A_i\rbrace _{i=1}^{N+1}$ for each object and the background are used as a vital cue to a segmentation head for the final segmentation. 3.2.3 Dynamic segmentation head We use a lightweight dynamic segmentation head that is dynamically instantiated once for each annotated object in the image with shared weights to efficiently segment the objects of interest from the background. Figure 4 presents the structure of the DSH. It consists of three depthwise separable convolution layers with a kernel size of 3 × 3 and widths of 256, 128 and 1. The input of the DSH for object i is a combination E i ∈ R w × h × 128 $E_i \in \mathbb {R}^{w\times h \times 128}$ of the appearance embeddings e, which are extracted from the image by the VFE, and the attention map A i $A_i$ of the object, which is computed from the user interactions by the RAM. Specifically, we first apply the attention map to each channel of the appearance embeddings by multiplication operation, and then perform a feature enhancement through accumulating the appearance embeddings. The combined features E i $E_i$ is computed as follows: E i c = A i ⊗ e c + e c = ( 1 + A i ) ⊗ e c , \begin{align} E_i^c = A_i \otimes e^c + e^c =(1+A_i)\otimes e^c, \end{align} (6)where E i c $E_i^c$ and e c $e^c$ indicate the feature slices of E i $E_i$ and e, respectively. ⊗ represents element-wise product. The residual formulation defined in (6) avoids missing meaningful information, while enhances object-relevant features efficiently. The output y i = D S H ( E i ) $y_i = DSH(E_i)$ of the DSH is a w × h × 1 $w\times h\times 1$ feature map of logits. FIGURE 4Open in figure viewerPowerPoint Architecture of the dynamic segmentation head. The lightweight segmentation head takes the features fused by distance maps and appearance embeddings and produces a one-dimensional feature map of logits for each object and the background. The logits are then stacked together and softmax is applied After the logits { y i } i N + 1 $\lbrace y_i\rbrace _i^{N+1}$ for each object and the background are extracted, we stack them together and apply a softmax layer to predict the objects and background (the right part of Figure 3): { M i } i N + 1 = f s o f t m a x ( { y i } i N + 1 ) , \begin{align} \lbrace M_i\rbrace _i^{N+1} = f_{softmax}(\lbrace y_i\rbrace _i^{N+1}), \end{align} (7)where f s o f t m a x ( · ) $f_{so