SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mengyu Wang $^{1,3*}$ , Henghui Ding $^{4}$ , Jun Hao Liew $^{5}$ , Jiajun Liu $^{5}$ , Yao Zhao $^{1,2,3\dagger}$ , Yunchao Wei $^{1,2,3\dagger}$
$^{4}$  Nanyang Technological University  $^{5}$  ByteDance
wmengyu0826@gmail.com
In this paper, we explore a principal way to enhance the quality of object masks produced by different segmentation models. We propose a model-agnostic solution called SegRefiner, which offers a novel perspective on this problem by interpreting segmentation refinement as a data generation process. As a result, the refinement process can be smoothly implemented through a series of denoising diffusion steps. Specifically, SegRefiner takes coarse masks as inputs and refines them using a discrete diffusion process. By predicting the label and corresponding states-transition probabilities for each pixel, SegRefiner progressively refines the noisy masks in a conditional denoising manner. To assess the effectiveness of SegRefiner, we conduct comprehensive experiments on various segmentation tasks, including semantic segmentation, instance segmentation, and dichotomous image segmentation. The results demonstrate the superiority of our SegRefiner from multiple aspects. Firstly, it consistently improves both the segmentation metrics and boundary metrics across different types of coarse masks. Secondly, it outperforms previous model-agnostic refinement methods by a significant margin. Lastly, it exhibits a strong capability to capture extremely fine details when refining high-resolution images. The source code and trained models are available at github.com/MengyuWang826/SegRefiner.
Although segmentation in image [35, 23, 24, 14, 11] and video [54, 37, 38, 16, 15] has been extensively studied in the past decades, obtaining accurate and detailed segmentation masks is always challenging since high-quality segmentation requires the model to capture both high-level semantic information and fine-grained texture information to make accurate predictions. This challenge is particularly pronounced for images with resolutions of  $2K$  or higher, which requires considerable computational complexity and memory usage. As a result, existing segmentation algorithms often predict masks at a smaller size, inevitably leading to lower accuracy due to the loss of fine-grained information during downsampling.
Since directly predicting high-quality masks is challenging, some previous works have shifted their attention to the refinement of coarse masks obtained from a preceding segmentation model. A popular line of direction [29, 57, 27, 28] is to augment the segmentation models (and features) with a new
*Work done during an internship at ByteDance.
† Co-corresponding author.
(a) Image
(b) Coarse Mask
(c) CRM [44]
Figure 1: Diverse errors presented in the previous segmentation results. The top row and middle row show the significant false positives and false negatives caused by incorrect semantics, the bottom row shows the inaccurate in capturing fine-grained details. Please zoom in for a better view.
(d) SegRefiner (ours)
module for masks correction. However, such approaches are usually model-specific, and hence cannot be generalized to refine coarse masks produced by other segmentation models. Another part segmentation refinement works [47, 56, 12, 44], on the other hand, resort to model-agnostic approach by taking only an image and the coarse segmentation as input for refinement. These methods have greater practical utility, as they are applicable to refine different segmentation models. However, the diverse types of errors (e.g., errors along object boundaries, failure in capturing fine-grained details in high-resolution images and errors due to incorrect semantics) presented in the coarse masks pose a great challenge to the refinement model thus causing underperformance (refer to Fig. 1).
In response to this challenge, we draw inspiration from the working principle of denoising diffusion models [45, 22, 2]. Diffusion models perform denoising at each timestep and gradually approach the image distribution through multiple iterations. This iterative strategy significantly reduces the difficulty of fitting the target distribution all at once, empowering diffusion models to generate high-quality images. Intuitively, by applying this strategy to the segmentation refinement task, refinement model can focus on correcting only some "most obvious errors" at each step and iteratively converge to an accurate result, thus reducing the difficulty of correcting all errors in a single pass and enabling refinement model to handle more challenging instances.
Under this perspective, we present an innovative interpretation of the task by representing segmentation refinement as a data generation process. As a result, refinement can be implemented through a sequence of denoising diffusion steps with coarse segmentation masks being the noisy version of ground truth. To work with binary masks, we further devise a novel discrete diffusion process, where every pixel performs unidirectional randomly states-transition. The proposed process can gradually convert ground truth into a coarse mask during training and employ the coarse mask as sampling origination during inference. In other words, we formulate the mask refinement task as a conditional generation problem, where the input image serves as the condition for iteratively updating/refining the erroneous predictions in the coarse mask.
To our best knowledge, we are the first to introduce diffusion-based refinement for segmentation masks. Our method, called SegRefiner, is model-agnostic, and thus applicable across different segmentation models and tasks. We extensively analyze the performance of SegRefiner across various segmentation tasks, demonstrating that our SegRefiner not only outperforms all previous model-agnostic refinement methods (with  $+3.42$  IoU,  $+2.21$  mBA in semantic segmentation and  $+0.9$  Mask AP,  $+2.2$  Boundary AP in instance segmentation), but can also be effortlessly transferred to other segmentation tasks (e.g., the recently proposed dichotomous image segmentation task [40]) without any modification. Additionally, SegRefiner demonstrates a strong capability for capturing extremely fine details when applied to high-resolution images (see Fig. 1 and Fig. 4).
Segmentation Refinement The aim of segmentation refinement is to improve the quality of masks in pre-existing segmentation models. Some works focus on enhancing specific segmentation models. PointRend [29] employs an MLP to predict the labels of pixels with low-confidence scores output from Mask R-CNN [21]. RefineMask [57] incorporates a semantic head to Mask R-CNN as additional guidance. MaskTransfiner [27] employs an independent FCN [35] to detect incoherent regions and refines their labels with a Transformer [49]. These works have demonstrated notable performance enhancements over their preceding models. However, their scope of improvement is model-specific and they lack the capacity to directly refine the coarse mask derived from other models. There are also some refinement methods that adopt model-agnostic approaches, such as [47, 56, 61, 12, 44, 31]. These strategies emphasize utilizing diverse forms of input, including whole images, boundary patches, edge strips, etc. Even though these techniques can refine coarse masks derived from different models, their applicability remains confined to specific segmentation tasks. As a special case, SegFix [56], which learns a mapping function between edge pixels and inner pixels and subsequently replaces inaccurate edge predictions with corresponding inner pixel predictions, is employed in both semantic segmentation and instance segmentation within the Cityscapes dataset [13]. While the performance of SegFix is significantly constrained by its ability to accurately identify objects within an image, which consequently leads to a decline in performance on datasets with a more extensive range of categories (e.g., COCO [34]).
Diffusion Models for Detection and Segmentation Recently, diffusion models have received a lot of attention in research. Initial studies [45, 22, 2, 46, 8, 3] primarily sought to enhance and expand the diffusion framework. Following these, subsequent works ventured to incorporate diffusion models across a broader array of tasks [43, 9, 30] and to formulate comprehensive conditional generation frameworks [41, 58]. The application of diffusion models to detection and segmentation tasks has also been the focus of an escalating number of studies. Baranchuk et al. [4] capture the intermediate activations from diffusion models and employ an MLP to execute per-pixel classification. DiffusionDet [6] and DiffusionInst [19] adapt the diffusion process to perform denoising in object boxes and mask filters. Some other works [1, 52, 53, 7, 26] consider the image segmentation task as mask generation. These studies predominantly follow the Gaussian diffusion process of DDPM [22] and leverage an additional image encoder to extract image features as condition to generate masks. To the best of our knowledge, our SegRefiner is the first work that applies a diffusion model to the image segmentation refinement task, and it also pioneers the abandonment of the Gaussian assumption in favor of a newly designed discrete diffusion process in diffusion-based segmentation tasks.
3.1 Preliminaries: Diffusion Models
Diffusion models consist of a forward and a reverse process. The forward process  $q(x_{1:T}|x_0)$  uses a Markov or None-Markov chain to gradually convert the data distribution  $x_0 \sim q(x_0)$  into complete noise  $x_T$  whereas the reverse process deploys a gradual denoising procedure  $p_\theta(x_{0:T})$  that transforms a random noise back into the original data distribution.
Continuous Diffusion Models The majority of existing continuous diffusion models [22, 46, 41] adheres to the Gaussian assumption and defines  $p(x_{T}) = \mathcal{N}(x_{T}|0,1)$ . The mean and variance of forward process are defined by a hyperparameter  $\beta_{t}$  and the reverse process utilizes mean and variance from model predictions, thus formulating as:
$$
q \left(x _ {t} \mid x _ {t - 1}\right) = \mathcal {N} \left(x _ {t} \mid \sqrt {1 - \beta_ {t}} x _ {t - 1}, \beta_ {t} I\right), \tag {1}
$$
$$
p _ {\theta} \left(x _ {t - 1} \mid x _ {t}\right) = \mathcal {N} \left(x _ {t - 1} \mid \mu_ {\theta} \left(x _ {t}, t\right), \Sigma_ {\theta} \left(x _ {t}, t\right)\right). \tag {2}
$$
Discrete Diffusion Models Compared to continuous diffusion models, there is less research on discrete diffusion models. Sohl-Dickstein et al. [45] first introduce binary diffusion to reconstruct one-dimensional noisy binary sequences.  $x_{T}$  is defined to adhere to the Bernoulli distribution  $\mathcal{B}(x_{T}|0.5)$ . The forward process and reverse process are represented as:
$$
q \left(x _ {t} \mid x _ {t - 1}\right) = \mathcal {B} \left(x _ {t} \mid x _ {t - 1} \left(1 - \beta_ {t}\right) + 0. 5 \beta_ {t}\right), \tag {3}
$$
$$
p _ {\theta} \left(x _ {t - 1} \mid x _ {t}\right) = \mathcal {B} \left(x _ {t - 1} \mid f _ {b} \left(x _ {t}, t\right)\right). \tag {4}
$$
Transition Sample
input as previous mask input as target mask
input as probability
element-wise addition
element-wise multiplication
Figure 2: An overview of the proposed SegRefiner (best viewed in color). On the right is the transition sample module we proposed, which randomly samples pixels from the current mask based on the input states-transition probabilities and change their values to match those in the target mask. During training, the transition sample module transforms the ground truth into a coarse mask, thus coarse mask is the target mask. During inference, target mask refers to the predicted fine mask and this module updates the values in the coarse mask in each timestep based on the predicted fine mask and the transition probabilities.
channel-dim concatenation
where  $\beta_{t} \in (0,1)$  is a hyperparameter and  $f_{b}(x_{t},t)$  is a model predicting Bernoulli probability. After the great success of DDPM [22], Austin et al. [2] extended the architecture of Discrete Diffusion Model to a more general form. They define the forward process as a discrete random variable transitioning among multiple states and use states-transition matrix  $Q_{t}$  to characterize this process:
$$
[ Q _ {t} ] _ {m, n} = q \left(x _ {t} = n \mid x _ {t - 1} = m\right). \tag {5}
$$
3.2 SegRefiner
In this work, we propose SegRefiner, with a unique discrete diffusion process, which can be applied to refine coarse masks from various segmentation models and tasks. SegRefiner performs refinement with a coarse-to-fine diffusion process. In the forward process, SegRefiner employs a discrete diffusion process which is formulated as unidirectional random states-transition, gradually degrading the ground truth mask into a coarse mask. In the reverse process, SegRefiner begins with a provided coarse mask and gradually transforms the pixels in the coarse mask to the refined state, correcting the wrongly predicted area in the coarse mask. In the following paragraphs, we will provide a detailed description of the forward and reverse process.
Forward diffusion process In the forward process, we gradually degrade the ground truth mask/ fine mask  $M_{fine}$ , transiting it into a coarse mask  $M_{coarse}$ . In other words, we have  $m_0 = M_{fine}$  and  $m_T = M_{coarse}$ . At any intermediate timestep  $t \in \{1, 2, \dots, T - 1\}$ , the intermediate mask  $m_t$  is therefore in a transitional phase between  $M_{fine}$  and  $M_{coarse}$ .
We define that every pixel in  $m_t$  has two states: fine and coarse, and the forward process is thus formulated as states-transition between these two states. Pixels in the fine state will retain their values from  $M_{fine}$ , and vice versa. We propose a new transition sample module to formulate this process. As shown in Fig. 2 right, during forward process, the transition sample module takes the previous mask  $m_{t-1}$ , coarse mask  $m_T$  and a states-transition probability as input and outputs a transitioned mask  $m_t$ . The states-transition probability describes the probability of every pixel in  $m_{t-1}$  transitioning to the coarse state. This module first performs Gumbel-max sampling [25] according to the states-transition probability and obtains the transitioned pixels. Then, the transitioned pixels will take values from  $m_T$  while the non-transitioned pixels will keep unchanged.
Note that the transition sample module represents a unidirectional process in which only "transition to coarse state" happens. The unidirectional property ensures that the forward process will converge
Figure 3: Examples of SegRefiner's inference process.  $x_{t}$  and  $m_{t}$  denote the state (we represent state [0, 1] and [1, 0] as 0 and 1 in this figure) and the corresponding mask, respectively. SegRefiner begins with  $x_{T} = [0, 1]$  and  $m_{T} = M_{\text{coarse}}$ , and gradually refines until we obtain fine mask  $m_{0}$ .
to  $M_{\text{coarse}}$ , despite each step is completely random. This is a significant difference between our SegRefiner and previous diffusion models in which forward process converges to a random noise.
With the reparameterization trick, we introduce a binary random variable  $x$  to formulate the above process. We represent  $x_{t}^{i,j}$  as a one-hot vector to represent the state of pixel  $(i,j)$  in  $m_{t}$  and set  $x_{0}^{i,j} = [1,0]$  and  $x_{T}^{i,j} = [0,1]$  to represent the fine state and coarse state, respectively. The forward process can thus be formulated as:
$$
q \left(x _ {t} ^ {i, j} \mid x _ {t - 1} ^ {i, j}\right) = x _ {t - 1} ^ {i, j} Q _ {t}, \quad \text {w h e r e} Q _ {t} = \left[ \begin{array}{c c} \beta_ {t} & 1 - \beta_ {t} \\ 0 & 1 \end{array} \right], \tag {6}
$$
where  $\beta_{t} \in [0,1]$ , and  $1 - \beta_{t}$  corresponds to the states-transition probability used in our transition sample module.  $Q_{t}$  is a states-transition matrix. The form of  $Q_{t}$  explicitly manifests the unidirectional property, i.e., all pixels in the coarse state will never transition back to the fine state since  $q(x_{t}|[0,1]) = [0,1]$ . According to Eq. (6), the marginal distribution can be formulated as
$$
q \left(x _ {t} ^ {i, j} \mid x _ {0} ^ {i, j}\right) = x _ {0} ^ {i, j} Q _ {1} Q _ {2} \dots Q _ {t} = x _ {0} \bar {Q} _ {t} = x _ {0} \left[ \begin{array}{c c} \bar {\beta} _ {t} & 1 - \bar {\beta} _ {t} \\ 0 & 1 \end{array} \right], \tag {7}
$$
where  $\bar{\beta}_t = \beta_1\beta_2\ldots \beta_t$ . Given this, we can obtain the intermediate mask  $m_{t}$  at any intermediate timestep  $t$  without the need of step-by-step sampling  $q(x_{t}|x_{t - 1})$ , allowing faster training.
Reverse diffusion process The reverse diffusion process takes a coarse mask  $m_T$  and gradually transforms it into a fine mask  $m_0$ . However, since the fine mask  $m_0$  and the reversed states-transition probability is unknown, following DDPM [22], we train a neural network  $f_{\theta}$  parameterized by  $\theta$  to predict the fine mask  $\tilde{m}_{0|t}$  at each timestep, represented as
$$
\tilde {m} _ {0 | t}, p _ {\theta} (\tilde {m} _ {0 | t}) = f _ {\theta} (I, m _ {t}, t), \tag {8}
$$
where  $I$  is the corresponding image.  $\tilde{m}_{0|t}$  and  $p_{\theta}(\tilde{m}_{0|t})$  denote the predicted binary fine mask and its confidence score, respectively. To obtain the reversed states-transition probability, according to Eq. (6), Eq. (7), and Bayes' theorem, we first formulate the posterior at timestep  $t - 1$  as:
$$
q \left(x _ {t - 1} \mid x _ {t}, x _ {0}\right) = \frac {q \left(x _ {t} \mid x _ {t - 1} , x _ {0}\right) q \left(x _ {t - 1} \mid x _ {0}\right)}{q \left(x _ {t} \mid x _ {0}\right)} = \frac {x _ {t} Q _ {t} ^ {\top} \odot x _ {0} \bar {Q} _ {t - 1}}{x _ {0} \bar {Q} _ {t} x _ {t} ^ {\top}}, \tag {9}
$$
where the fine state  $x_0$  is set to  $[1, 0]$  during training, indicating ground truth. While during inference,  $x_0$  is unknown, as the predicted  $\tilde{m}_{0|t}$  may not be entirely accurate. Since the confidence score  $p_{\theta}(\tilde{m}_{0|t})$  represents the model's level of certainty for each pixel prediction being correct,  $p_{\theta}(\tilde{m}_{0|t})$  can also be interpreted as the probability being in the fine state. Therefore, intuitively, one could
obtain the state of every pixel in  $\tilde{m}_{0|t}$  by simply thresholding as done in [2]:
$$
x _ {0 \mid t} ^ {i, j} = \left\{ \begin{array}{l l} [ 1, 0 ] & \text {i f} p _ {\theta} \left(\tilde {m} _ {0 \mid t}\right) ^ {i, j} \geq 0. 5 \\ [ 0, 1 ] & \text {o t h e r w i s e}, \end{array} \right. \tag {10}
$$
where pixels with higher confidence scores will have  $x_{0|t}^{i,j} = [1,0]$ , indicating they are in the fine state, and vice versa. However, in such one-hot form, the values of states-transition probabilities are determined solely by the pre-defined hyperparameters, leading to significant information loss. Instead, we retain the soft transition and formulate  $x_{0|t}^{i,j} = [p_{\theta}(\tilde{m}_{0|t})^{i,j}, 1 - p_{\theta}(\tilde{m}_{0|t})^{i,j}]$ . With this setting, the reverse diffusion process can be reformulated as
$$
p _ {\theta} \left(x _ {t - 1} ^ {i, j} \mid x _ {t} ^ {i, j}\right) = x _ {t} ^ {i, j} P _ {\theta , t} ^ {i, j}, \quad \text {w h e r e} P _ {\theta , t} ^ {i, j} = \left[ \begin{array}{c c} 1 & 0 \\ \frac {p _ {\theta} \left(\tilde {m} _ {0 , t}\right) ^ {i , j} \left(\bar {\beta} _ {t - 1} - \bar {\beta} _ {t}\right)}{1 - p _ {\theta} \left(\tilde {m} _ {0 , t}\right) ^ {i , j} \bar {\beta} _ {t}} & \frac {1 - p _ {\theta} \left(\tilde {m} _ {0 , t}\right) ^ {i , j} \bar {\beta} _ {t - 1}}{1 - p _ {\theta} \left(\tilde {m} _ {0 , t}\right) ^ {i , j} \bar {\beta} _ {t}} \end{array} \right], \tag {11}
$$
where  $P_{\theta, t}^{i,j}$  is the reversed states-transition matrix. With the above reversed states-transition probability,  $m_t$  and  $\tilde{m}_{0|t}$  as input, the transition sample module can transit a portion of pixels to the fine state at each timestep, thereby correcting erroneous predictions.
Inference Given a coarse mask  $m_T$  and its corresponding image  $I$ , we first initialize that all pixels are in the coarse state thus  $x_{T}^{i,j} = [0,1]$ . We iterate between: (1) forward pass to obtain  $\tilde{m}_{0|t}$  and  $p_{\theta}(\tilde{m}_0|t)$  (Eq. (8)); (2) compute the reversed states-transition matrix  $P_{\theta ,t}^{i,j}$  and obtain  $x_{t - 1}$  (Eq. (11)); (3) compute the refined mask  $m_{t - 1}$  based on  $x_{t - 1}, m_t$  and  $\tilde{m}_{0|t}$ . The process (1)-(3) is iterated until we obtain fine mask  $m_0$ . Visualization examples of inference are shown in Fig. 3.
Model Architecture Following [39], we employ U-Net for our denoising network. We modify the U-Net to take in 4-channel input (concatenation of image and the corresponding mask  $m_{t}$ ) and output a 1-channel refined mask. Both input and output resolution is set to  $256 \times 256$ . All others remain unchanged other than the aforementioned modifications.
Objective Function Following [12], we employ a combination of binary cross-entropy loss and texture loss for training our model, i.e.,  $\mathcal{L} = \mathcal{L}_{bce} + \alpha \mathcal{L}_{texture}$ , where texture loss is characterized as an L1 loss between the segmentation gradient magnitudes of the predicted mask and the ground truth mask.  $\alpha$  is set to 5 to balance the magnitude of both losses.
Noise Schedule Theoretically, the unidirectional property of SegRefiner ensures that any noise schedule can make the forward process converge to the coarse mask given an infinite number of timesteps. However, in practice, we use much fewer timesteps ( $T = 6$  this work) to ensure efficient inference. We designate  $\bar{\beta}_T = 0$  such that  $x_T = [0,1]$  for all pixels and  $m_T = M_{\text{coarse}}$  (Eq. (7)). Following DDIM [46], we directly set a linear noise schedule from 0.8 to 0 on  $\bar{\beta}_t$ .
Training Strategy Our SegRefiner model has been developed into two versions for refinement of different resolution images: a low-resolution variant (hereafter referred to as LR-SegRefiner) and a high-resolution variant (hereafter referred to as HR-SegRefiner). While these two versions employ different training datasets, all other settings remain consistent. LR-SegRefiner is trained on the LVIS dataset [20], whereas HR-SegRefiner is trained on a composite dataset merged from two high-resolution datasets, DIS5K [40] and ThinObject-5K [32]. These datasets were chosen due to their highly accurate pixel-level annotations, thereby facilitating the training of our model to capture fine details more effectively. Following [12], the coarse masks used for training are obtained through various morphological operations, such as randomly perturbing some edge points of the ground truth and performing dilation, erosion, etc. During training, we first train LR-SegRefiner on the low-resolution dataset until convergence. Subsequently, it is fine-tuned on the high-resolution dataset to yield HR-SegRefiner.
We employed double random crop as the primary data augmentation technique. It entails the random cropping of each image to generate two distinct model inputs: one encapsulating the entire foreground
target, and the other containing only a partial foreground. Both crops are subsequently resized to match the model's input size. This operation ensures that the model is proficient in refining both entire objects and incomplete local patches, a capability that proved instrumental in subsequent experiments. All the following experiments were conducted on 8 NVIDIA RTX 3090. For more details on the training process, please refer to the supplemental materials.