SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Wentong  $\mathbf{L}\mathbf{i}^{1*}$  , Yuqian Yuan $^{1*}$ , Song Wang $^{1}$ , Wenyu Liu $^{1}$ ,
Dongqi Tang $^{2}$ , Jian Liu $^{2}$ , Jianke Zhu $^{1\dagger}$ , Lei Zhang $^{3}$
$^{1}$ Zhejiang University  $^{2}$ Ant Group  $^{3}$ The Hong Kong Polytechnical University
https://LiWentomng.github.io/apro/
Weakly-supervised segmentation with label-efficient sparse annotations has attracted increasing research attention to reduce the cost of laborious pixel-wise labeling process, while the pairwise affinity modeling techniques play an essential role in this task. Most of the existing approaches focus on using the local appearance kernel to model the neighboring pairwise potentials. However, such a local operation fails to capture the long-range dependencies and ignores the topology of objects. In this work, we formulate the affinity modeling as an affinity propagation process, and propose a local and a global pairwise affinity terms to generate accurate soft pseudo labels. An efficient algorithm is also developed to reduce significantly the computational cost. The proposed approach can be conveniently plugged into existing segmentation networks. Experiments on three typical label-efficient segmentation tasks, i.e. box-supervised instance segmentation, point/scribble-supervised semantic segmentation and CLIP-guided semantic segmentation, demonstrate the superior performance of the proposed approach.
Segmentation is a widely studied problem in computer vision, aiming at generating a mask prediction for a given image, e.g., grouping each pixel to an object instance (instance segmentation) or assigning each pixel a category label (semantic segmentation). While having achieved promising performance, most of the existing approaches are trained in a fully supervised manner, which heavily depend on the pixel-wise mask annotations, incurring tedious labeling costs [1]. Weakly-supervised methods have been proposed to reduce the dependency on dense pixel-wise labels with label-efficient sparse annotations, such as points [2-4], scribbles [5-7], bounding boxes [8-11] and image-level labels [12-15]. Such methods make dense segmentation more accessible with less annotation costs for new categories or scene types.
Most of the existing weakly-supervised segmentation methods [16, 13, 8, 17, 3, 10] adopt the local appearance kernel to model the neighboring pairwise affinities, where spatially nearby pixels with similar color (i.e., LAB color space [8, 3] or RGB color space [13, 16, 17, 10]) are likely to be in the same class. Though having proved to be effective, these methods suffer from two main limitations. First, the local operation cannot capture global context cues and capture long-range affinity dependencies. Second, the appearance kernel fails to take the intrinsic topology of objects into account, and lacks capability of detail preservation.
To address the first issue, one can directly enlarge the kernel size to obtain a large receptive filed. However, this will make the segmentation model insensitive to local details and increase the computational cost greatly. Some methods [14, 12] model the long-range affinity via random walk [18], but they cannot model the fine-grained semantic affinities. As for the second issue, the tree-based
† Correspondence author
Figure 1: The proposed approach upon the typical weakly-supervised segmentation tasks with label-efficient annotations, including (a) box-supervised instance segmentation, (b) point/scribble-supervised semantic segmentation and (c) annotation-free semantic segmentation with CLIP pretrained model.
approaches [7, 19] are able to preserve the geometric structures of objects, and employ the minimum spanning tree [20] to capture the pairwise relationship. However, the affinity interactions with distant nodes will decay rapidly as the distance increases along the spanning tree, which still focuses on the local nearby regions. LTF-V2 [21] enables the long-range tree-based interactions but it fails to model the valid pairwise affinities for label-efficient segmentation task.
With the above considerations, we propose a novel component, named Affinity Propagation (APro), which can be easily embedded in existing methods for label-efficient segmentation. Firstly, we define the weakly-supervised segmentation from a new perspective, and formulate it as a uniform affinity propagation process. The modelled pairwise term propagates the unary term to other nearby and distant pixels and updates the soft pseudo labels progressively. Then, we introduce the global affinity propagation, which leverages the topology-aware tree-based graph and relaxes the geometric constraints of spanning tree to capture the long-range pairwise affinity. With the efficient design, the  $\mathcal{O}(N^2)$  complexity of brute force implementation is reduced to  $\mathcal{O}(N\log N)$ , and the global propagation approach can be performed with much less resource consumption for practical applications. Although the long-range pairwise affinity is captured, it inevitably brings in noises based on numerous pixels in a global view. To this end, we introduce a local affinity propagation to encourage the piece-wise smoothness with spatial consistency. The formulated APro can be embedded into the existing segmentation networks to generate accurate soft pseudo labels online for unlabeled regions. As shown in Fig. 1, it can be seamlessly plugged into the existing segmentation networks for various tasks to achieve weakly-supervised segmentation with label-efficient sparse annotations.
We perform experiments on three typical label-efficient segmentation tasks, i.e. box-supervised instance segmentation, point/scribble-supervised semantic segmentation and annotation-free semantic segmentation with pretrained CLIP model, and the results demonstrated the superior performance of our proposed universal label-efficient approach.
Label-efficient Segmentation. Label-efficient segmentation, which is based on the weak supervision from partial or sparse labels, has been widely explored [1]. Different from semi-supervised settings [22, 23], this paper mainly focuses on the segmentation with sparse labels. In earlier literature [24-28], it primarily pertained to image-level labels. Recently, diverse sparse annotations have been employed, including point, scribble, bounding box, image-level label and the combinations of them. We briefly review the weakly-supervised instance segmentation, semantic segmentation and panoptic segmentation tasks in the following.
For weakly-supervised instance segmentation, box-supervised methods [16, 8, 17, 9, 29, 30, 11, 10, 31] are dominant and perform on par with fully-supervised segmentation approaches. Besides, the "points + bounding box" annotation can also achieve competitive performance [32, 33]. As for weakly-supervised semantic segmentation, previous works mainly focus on the point-level supervision [2, 34, 6] and scribble-level supervision [5, 35, 36], which utilize the spatial and color information of the input image and are trained with two stages. Liang et al. [7] introduced an effective tree energy loss based on the low-level and high-level features for point/scribble/block-supervised semantic segmentation. For semantic segmentation, the supervision of image-level labels has been well explored [12, 37, 15, 14]. Recently, some works have been proposed to make use of the large-scale pretrained CLIP model [38] to achieve weakly-supervised or annotation-free semantic
(b)
Figure 2: Overview of our APro approach. (a) The general weakly supervised segmentation framework with the proposed pairwise affinity propagation. (b) The proposed approach consists of global affinity propagation (GP) and local affinity propagation (LP) to generate accurate pseudo labels.
segmentation [39, 40]. In addition, weakly-supervised panoptic segmentation methods [3, 41, 4] with a single and multiple points annotation have been proposed. In this paper, we aim to develop a universal component for various segmentation tasks, which can be easily plugged into the existing segmentation networks.
Pairwise Affinity Modeling. Pairwise affinity modeling plays an important role in many computer vision tasks. Classical methods, like CRF [42], make full use of the color and spatial information to model the pairwise relations in the semantic labeling space. Some works use it as a post-processing module [43], while others integrate it as a jointly-trained part into the deep neural network [44, 45]. Inspired by CRF, some recent approaches explore the local appearance kernel to tap the neighboring pairwise affinities, where spatially nearby pixels with similar colors are more likely to fall into the same class. Tian et al. [8] proposed the local pairwise loss, which models the neighboring relationship in LAB color space. Some works [10, 16, 17] focus on the pixel relation in RGB color space directly, and achieve competitive performance. However, the local operation fails to capture global context cues and lacks the long-range affinity dependencies. Furthermore, the appearance kernel cannot reflect the topology of objects, missing the details of semantic objects. To model the structural pair-wise relationship, tree-based approaches [7, 19] leverage the minimum spanning tree [20] to capture the topology of objects in the image. However, the affinity interactions with distant nodes decay rapidly as the distance increases along the spanning tree. Besides, Zhang et al. [35] proposed an affinity network to convert an image to a weighted graph, and model the node affinities by the graph neural network (GNN). Different from these methods, in this work we model the pairwise affinity from a new affinity propagation perspective both globally and locally.
In this section, we introduce our proposed affinity propagation (APro) approach to label-efficient segmentation. First, we define the problem and model the APro framework in Section 3.1. Then, we describe the detailed pairwise affinity propagation method in Section 3.2. Finally, we provide an efficient implementation of our method in Section 3.3.
3.1 Problem Definition and Modeling
Given an input image  $I = \{x_{i}\}^{N}$  with  $N$  pixels and its available sparse ground truth labeling  $\bar{Y}$  (i.e. points, scribble, or bounding box), and let  $f_{\theta}(I)$  be the output of a segmentation network with the learnable parameters  $\theta$ , the whole segmentation network can be regarded as a neural network optimization problem as follows:
$$
\min  _ {\theta} \left\{\mathcal {L} _ {g} \left(f _ {\theta} (I), \bar {Y}\right) + \mathcal {L} _ {u} \left(f _ {\theta} (I), Y\right) \right\}, \tag {1}
$$
where  $\mathcal{L}_g$  is a ground truth loss on the set of labeled pixels  $\Omega_g$  with  $\bar{Y}$  and  $\mathcal{L}_u$  is a loss on the unlabeled regions  $\Omega_u$  with the pseudo label  $Y$ . As shown in Fig. 2-(a), our goal is to obtain the accurate pseudo label  $Y$  by leveraging the image pixel affinities for unlabeled regions.
As in the classic MRF/CRF model [46, 42], the unary term reflects the per-pixel confidence of assigning labels, while pairwise term captures the inter-pixel constraints. We define the generation of
pseudo label  $Y$  as an affinity propagation process, which can be formulated as follows:
$$
y _ {i} = \frac {1}{z _ {i}} \sum_ {j \in \tau} \phi \left(x _ {j}\right) \psi \left(x _ {i}, x _ {j}\right), \tag {2}
$$
where  $\phi(x_j)$  denotes the unary potential term, which is used to align  $y_j$  and the corresponding network prediction based on the available sparse labels.  $\psi(x_i, x_j)$  indicates the pairwise potential, which models the inter-pixel relationships to constrain the predictions and produce accurate pseudo labels.  $\tau$  is the region with different receptive fields.  $z_i$  is the summation of pairwise affinity  $\psi(x_i, x_j)$  along with  $j$  to normalize the response.
Notably, we unify both global and local pairwise potentials in an affinity propagation process formulated in Eq. 2. As shown in Fig. 2-(b), the global affinity propagation (GP) can capture the pairwise affinities with topological consistency in a global view, while the local affinity propagation (LP) can obtain the pairwise affinities with local spatial consistency. Through the proposed component, the soft pseudo labels  $Y$  can be obtained. We assign each  $y_{i}$  from GP and LP with the network prediction  $p_i$  and directly employ the distance measurement function as the objective for unlabeled regions  $\Omega_u$ . Simple  $L_1$  distance is empirically adopted in our implementation.
3.2 Pairwise Affinity Propagation
3.2.1 Global Affinity Propagation
We firstly provide a solution to model the global affinity efficiently based on the input image. Specifically, we represent an input image as a 4-connected planar graph  $\mathcal{G}$ , where each node is adjacent to up to 4 neighbors. The weight of the edge measures the image pixel distance between adjacent nodes. Inspired by tree-based approaches [47, 19], we employ the minimum spanning tree (MST) algorithm [20] to remove the edge with a large distance to obtain the tree-based sparse graph  $\mathcal{G}_T$ , i.e.,  $\mathcal{G}_T \gets \mathrm{MST}(\mathcal{G})$  and  $\mathcal{G}_T = \{\mathcal{V}, \mathcal{E}\}$ , where  $\mathcal{V} = \{\mathcal{V}_i\}^N$  is the set of nodes and  $\mathcal{E} = \{\mathcal{E}_i\}^{N-1}$  denotes the set of edges.
Then, we model the global pairwise potential by iterating over each node. To be specific, we take the current node as the root of the spanning tree  $\mathcal{G}_T$  and propagate the long-range affinities to other nodes. While the distant nodes along the spanning tree need to pass through nearby nodes along the path of spanning tree, the distance-insensitive max affinity function can alleviate this geometric constraint and relax the affinity decay for long-range nodes. Hence, we define the global pairwise potential  $\psi_g$  as follows:
$$
\psi_ {g} \left(x _ {i}, x _ {j}\right) = \mathcal {T} \left(I _ {i}, I _ {j}\right) = \exp \left(- \max  _ {\forall j \in \mathcal {V}} \frac {w _ {k , l}}{\zeta_ {g} {} ^ {2}}\right), \tag {3}
$$
where  $\mathcal{T}$  denotes the global tree.  $\mathbb{E}_{i,j}$  is the set of edges in the path of  $\mathcal{T}$  from node  $j$  to node  $i$ .  $w_{k,l}$  indicates the edge weight between the adjacent nodes  $k$  and  $l$ , which is represented as the Euclidean distance between pixel values of two adjacent nodes, i.e.,  $w_{k,l} = |I_k - I_l|^2$ .  $\zeta_g$  controls the degree of similarity with the long-range pixels. In this way, the global affinity propagation (GP) process to obtain the soft label  $y^g$  can be formulated as follows:
$$
y _ {i} ^ {g} = \frac {1}{z _ {i} ^ {g}} \sum_ {j \in \mathcal {V}} \phi (x _ {j}) \psi_ {g} (x _ {i}, x _ {j}), z _ {i} ^ {g} = \sum_ {j \in \mathcal {V}} \psi_ {g} (x _ {i}, x _ {j}), \tag {4}
$$
where interactions with distant nodes are performed over tree-based topology. In addition to utilizing low-level image, we empirically employ high-level feature as input to propagate semantic affinity.
3.2.2 Local Affinity Propagation
The long-range pairwise affinity is inevitably noisy since it is computed based on the susceptible image features in a global view. The spatially nearby pixels are more likely to have the same label, while they have certain difference in color and intensity, etc. Hence, we further introduce the local affinity propagation (LP) to promote the piece-wise smoothness. The Gaussian kernel is widely used to capture the local relationship among the neighbouring pixels in previous works [45, 8, 10]. Different from these works, we define the local pairwise affinity via the formulated affinity propagation process. The local pairwise term  $\psi_{s}$  is defined as:
$$
\psi_ {s} \left(x _ {i}, x _ {j}\right) = \mathcal {K} \left(I _ {i}, I _ {j}\right) = \exp \left(\frac {- \left| I _ {i} - I _ {j} \right| ^ {2}}{\zeta_ {s} ^ {2}}\right), \tag {5}
$$
where  $\mathcal{K}$  denotes the Gaussian kernel,  $\mathcal{N}(i)$  is the set containing all local neighbor pixels. The degree of similarity is controlled by parameter  $\zeta_{s}$ . Then the pseudo label  $y^{s}$  can be obtained via the following affinity propagation:
$$
y _ {i} ^ {s} = \frac {1}{z _ {i} ^ {s}} \sum_ {j \in \mathcal {N} (i)} \phi \left(x _ {j}\right) \psi_ {s} \left(x _ {i}, x _ {j}\right), z _ {i} ^ {s} = \sum_ {j \in \mathcal {N} (i)} \psi_ {s} \left(x _ {i}, x _ {j}\right), \tag {6}
$$
where the local spatial consistency is maintained based on high-contrast neighbors. To obtain a robust segmentation performance, multiple iterations are required. Notably, our LP process ensures a fast convergence, which is  $5 \times$  faster than MeanField-based method [10, 42]. The details can be found in the experimental section 4.4.
3.3 Efficient Implementation
Given a tree-based graph  $\mathcal{G}_T = \{\mathcal{V},\mathcal{E}\}$  in the GP process, we define the maximum  $\boldsymbol{w}$  value of the path through any two vertices as the transmission cost  $\mathcal{C}$ . One straightforward approach to get  $y_{i}^{g}$  of vertex  $i$  is to traverse each vertex  $j$  by Depth First Search or Breadth First Search to get the transmission cost  $\mathcal{C}_{i,j}$  accordingly. Consequently, the computational complexity required to acquire the entire output  $y^{g}$  is  $\mathcal{O}(N^2)$ , making it prohibitive in real-world applications.
Instead of calculating and updating the transmission cost of any two vertices, we design a lazy update algorithm to accelerate the GP process. Initially, each node is treated as a distinct union, represented by  $\mathcal{U}$ . Unions are subsequently connected based on each edge  $w_{k,l}$  in ascending order of  $\boldsymbol{w}$ . We show that when connecting two unions  $\mathcal{U}_k$  and  $\mathcal{U}_l$ ,  $w_{k,l}$  is equivalent to the transmission cost for all nodes within  $\mathcal{U}_k$  and  $\mathcal{U}_l$ . This is proved in the supplementary material.
To efficiently update values, we introduce a Lazy Propagation scheme. We only update the value of the root node and postpone the update of its descendants. The update information is retained in a lazy tag  $\mathcal{Z}$  and is updated as follows:
$$
\mathcal {Z} (\delta) _ {k ^ {*}} = \mathcal {Z} (\delta) _ {k ^ {*}} + \left\{ \begin{array}{l l} \exp \left(- w _ {k, l} / \zeta_ {g} ^ {2}\right) S (\delta) _ {l} & \mathcal {U} _ {k}. \text {r a n k} > \mathcal {U} _ {l}. \text {r a n k}, \\ \exp \left(- w _ {k, l} / \zeta_ {g} ^ {2}\right) S (\delta) _ {l} - \mathcal {Z} (\delta) _ {l ^ {*}} & \text {o t h e r w i s e}, \end{array} \right. \tag {7}
$$
where  $S(\delta)_i = \sum_{j\in \mathcal{U}_i}\delta_j$ $\delta$  means different inputs, including the dense prediction  $\phi (x)$  and all-one matrix  $\Lambda$ $k^{*} / l^{*}$  denotes the root node of node  $k / l$
Once all unions are connected, the lazy tags can be propagated downward from the root node to its descendants. For the descendants, the global affinity propagation term is presented as follows:
$$
L P r o p (\delta) _ {i} = \delta_ {i} + \sum_ {r \in A s c \mathcal {G} _ {T} (i) \cup \{i \}} \mathcal {Z} (\delta) _ {r}, \tag {8}
$$
where  $\operatorname{Asc}_{\mathcal{G}_T}(i)$  represents the ascendants of node  $i$  in the tree  $\mathcal{G}_T$ . As shown in Algorithm 1, the disjoint-set data structure is employed to implement the proposed algorithm. In our implementation, a Path Compression strategy is applied, connecting each node on the path directly to the root node. Consequently, it is sufficient to consider the node itself and its parent node to obtain LProp.
Time complexity. For each channel of the input, the average time complexity of sorting is  $\mathcal{O}(N\log N)$ . In the merge step, we utilize the Path Compression and Union-by-Rank strategies, which have a complexity of  $\mathcal{O}(\alpha (N))[48]$ . After merging all the concatenated blocks, the lazy tags can be propagated in  $\mathcal{O}(N)$  time. Hence, the overall complexity is  $\mathcal{O}(N\log N)$ . Note that the batches and channels are independent of each other. Thus, the algorithm can be executed in parallel for both batches and channels for practical implementations. As a result, the proposed algorithm reduces the computational complexity dramatically.
4.1 Weakly-supervised Instance Segmentation
Datasets. As in prior arts [8, 9, 16, 17], we conduct experiments on two widely used datasets for the weakly box-supervised instance segmentation task:
COCO [49], which has 80 classes with 115K train2017 images and 5K val2017 images.
Algorithm 1: Algorithm for GP process
Input: Tree  $\mathcal{G}_T\in \mathbb{N}^{e\times 2}$  ; Pairwise distance  $\pmb {w}\in \mathbb{R}^N$  ; Dense predictions  $\phi (x)\in \mathbb{R}^N$  Vertex num  $N$  ; Edge num  $e = N - 1$  ; Set of vertices  $\nu$    
Output:  $y^{g}\in \mathbb{R}^{N}$ $\Lambda \gets 1\in \mathbb{R}^{N}$ $F\gets \{0,1,2,\dots,N - 1\}$  ▷ Initialize each vertex as a connected block   
Sort  $\{\mathcal{G}_T,w\}$  in ascending order of w. ▷ Quick Sort   
for  $(k,l)\in \mathcal{G}_T,w_i\in \pmb {w}$  do a← find(k), b← find(l) ▷ Find the root node with Path Compression Update{Z(φ)a,Z(Λ)a,Z(φ)b,Z(Λ)b} ▷ Add lazy tag if  $S_{a} <   S_{b}$  then swap(a,b) ▷ Merge by Rank Fb← a ▷ Merge two connected blocks   
for  $v\in \mathcal{V}$  do p← find(v) for  $\delta \in \{\phi ,\Lambda \}$  do if  $p = v$  then LProp(δ)v= Z(δ)v+δv else LProp(δ)v= Z(δ)p+Z(δ)v+δv  $y_{v}^{g} = \frac{LProp(\phi)_{v}}{LProp(\Lambda)_{v}}$  ▷ Normalization   
return yg
- Pascal VOC [43] augmented by SBD [50] based on the original Pascal VOC 2012 [51], which has 20 classes with 10,582 trainaug images and 1,449 val images.
Base Architectures and Competing Methods. In the evaluation, we apply our proposed APro to two representative instance segmentation architectures, SOLOv2 [52] and Mask2Former [53], with different backbones (i.e., ResNet [54], Swin-Transformer [55]) following Box2Mask [11]. We compare our approach with its counterparts that model the pairwise affinity based on the image pixels without modifying the base segmentation network for box-supervised setting. Specifically, the compared methods include Pairwise Loss [8], TreeEnergy Loss [7] and CRF Loss [10]. For fairness, we re-implement these models using the default setting in MMDetection [56].
Implementation Details. We follow the commonly used training settings on each dataset as in MMDetection [56]. All models are initialized with ImageNet [57] pretrained backbone. For SOLOv2 framework [52], the scale jitter is used, where the shorter image side is randomly sampled from 640 to 800 pixels. For Mask2Former framework [53], the large-scale jittering augmentation scheme [58] is employed with a random scale sampled within range [0.1, 2.0], followed by a fixed size crop to  $1024 \times 1024$ . The initial learning rate is set to  $10^{-4}$  and the weight decay is 0.05 with 16 images per mini-batch. The box projection loss [8, 9] is employed to constrain the network prediction within the bounding box label as the unary term  $\phi$ . COCO-style mask AP (\%) is adopted for evaluation.
Quantitative Results. Table 1 shows the quantitative results. We compare the approaches with the same architecture for fair comparison. The state-of-the-art methods are listed for reference. One can see that our APro method outperforms its counterparts across Pascal VOC and COCO datasets.