EyeSeg: An Uncertainty-Aware Eye Segmentation Framework for AR/VR

1 Introduction

In recent years, human-machine interaction, especially augmented reality (AR) and virtual reality (VR), is increasingly being popularized by major smart hardware products. For such products to deliver smooth interaction, it is essential to have a gaze estimation algorithm to capture users' iris location and movement accurately and efficiently. Specifically, the algorithm takes as input a real-time image captured by a camera installed in the equipment and returns the resulting segmentation image that specifies the location of the iris in a pixel-wise manner. This segmentation-based product solution usually confronts three challenges: (i) imagery blurriness: due to eye fast motion, images captured usually exhibit blurriness, which poses difficulty to accurate eye segmentation; (ii) eyelid occlusion; (iii) train-test domain gaps: users may use the equipment in various device configurations or varying usage habits, hence inevitably there exists domain gaps between training data and test data. We demonstrate these challenges in Fig. 4.

♯ Corresponding Author

Humans typically express their uncertainty when faced with blurry information or when there are biases in domain knowledge; this allows them to make improvements or reject uncertain conclusions in subsequent analyzes. Inspired by this human trait, we propose a novel uncertainty-aware framework for eye segmentation (EyeSeg) designed to resolve the aforementioned challenges. Unlike previous methods [Wei et al., 2021; Czolbe et al., 2021; Sirohi et al., 2023; Fuhl et al., 2024] that rely on indirect or proxy measures of uncertainty, our approach explicitly models the uncertainty of the segmentation results. Specifically, we address the variations of input iris images within the eye regions which form a closed set hereinafter. Compared to open-set, the nature of closed-set conveniently induces a desired prior which a posterior approximates to achieve Bayesian uncertainty learning. Leveraging this prior, we utilize a statistic of the learned posterior to estimate posterior probabilities. Theoretically, we prove that the statistic of the learned posterior reflects the level of uncertainty for segmentation results.

In our framework, beyond the conventional segmentation pipeline, we extract multi-level features from the segmentation network and predict per-pixel variance, which is then used to compute the overall uncertainty of the image. During the testing phase, if the uncertainty exceeds a predefined threshold, the corresponding prediction is discarded. The pipeline is shown in Figure 1. This uncertainty score can be used as a weight to yield a more reliable gaze estimate, especially in the conditions of eyelid occlusion, imagery blurriness and in-domain/cross-domain settings. Moreover, our algorithm operates directly on the detected eye patch, leading to significant improvements in metrics such as Mean Intersection over Union (MIoU), E1, F1, and ACC, while maintaining computational efficiency with only 1.53G FLOPs.

To sum up, our major contributions are listed ´Ias follows:

• (Novelty) We propose a novel uncertainty-aware eye segmentation framework that resolves common challenges in AR/VR.
• (Interpretability) We prove that a statistic of the learned posterior provably reflects the level of uncertainty of segmentation results.
• (Evaluation) We conduct extensive experiments on multiple real-world datasets and challenge situations, demonstrating significant improvements in metrics such as MIoU, E1, F1, and ACC with 1.53G FLOPs.

2 RelatedWork

Eye Segmentation. Semantic segmentation is a fundamental task in computer vision with numerous applications, and there have been significant advancements [Long et al., 2015; Ronneberger et al., 2015; Badrinarayanan et al., 2017; Chen et al., 2014; Chen et al., 2017a; Chen et al., 2017b; Chen et al., 2018; Gong et al., 2021; Xie et al., 2021; Yuan et al., 2023; Tan et al., 2023; Peng et al., 2023; Yuan et al., 2023; Erisen, 2024; Chen et al., 2024; Liu et al., 2024] in recent years. An important application scenario is gaze estimation in AR/VR, where eye segmentation serves as a crucial preprocessing step. Various approaches[Hansen and Pece, 2005; ´Swirski et al., 2012; Chaudhary et al., 2019; Yiu et al., 2019; Feng et al., 2022; Luo et al., 2020; Wang et al., 2020; Kothari et al., 2021; Biswas and Lescroart, 2023; Fuhl et al., 2024] have been developed, taking into account both the unique characteristics of eye images and the real-time requirements of such applications. Traditional methods[Hansen and Pece, 2005; ´Swirski et al., 2012] typically employ morphological operations and geometric features to extract pupil regions from eye images. Deep learning-based methods[Chaudhary et al., 2019; Yiu et al., 2019; Feng et al., 2022; Luo et al., 2020; Wang et al., 2020; Kothari et al., 2021; Biswas and Lescroart, 2023], particularly convolutional neural networks, have shown promising results in eye segmentation. These methods leverage the ability of CNNs to learn hierarchical features from large amounts of data, which can be beneficial in handling complex and challenging scenarios. Some CNNs methods[Wang et al., 2020; Kothari et al., 2021; Biswas and Lescroart, 2023] also utilize ellipse fitting to obtain the shapes of the iris and pupil. Unlike these methods mentioned above, our work develops an uncertainty-aware eye segmentation framework which provides theoretically-grounded uncertainty score besides eye segmentation results, offering an efficient, accurate, and robust solution for eye segmentation. Recently, transformer-based methods [Hassan et al., 2022; Wei et al., 2021] have been proposed for eye segmentation. However, due to their high computational complexity, these methods often fail to meet real-time requirements, making them unsuitable for AR/VR applications.

Uncertainty Estimation. Existing methods mainly includes Probabilistic methods [Kendall and Gal, 2017; Czolbe et al., 2021], Evidence Learning[Sirohi et al., 2023], Deterministic Methods [Czolbe et al., 2021; Wei et al., 2021] and Ensemble methods [Czolbe et al., 2021; Landgraf et al., 2024]. In particular, BiTrans [Wei et al., 2021] leverages uncertainty to refine the eye segmentation, but it is not designed for eye segmentation in VR/AR scenes. Our approach belongs to the category of probabilistic methods and exhibits a high degree of interpretability. In contrast to other probabilistic methods, our approach directly models uncertainty in a closed set, where uncertainty is computed only at the subregion level. We leverage uncertainty estimation to reject erroneous outputs, instead of using it solely for improving model performance.

3 Methodology

Our proposed approach, EyeSeg, first detects eye regions, followed by uncertainty-aware eye segmentation. Eye regions provide a closed set, which conveniently induces a desired prior for posterior approximation to achieve uncertainty awareness. The entire approach is illustrated in Fig. 1.

Notation. Throughout the paper, we let (X, Y ) be a training pair drawn from a data set, where X denote the raw grayscale training image and Y the corresponding K-class pixellevel label (i.e. the segmentation map). Note that X and Y are two-dimensional tensors of which Xi,j represents the pixel value and Yi,j denotes the segmentation label at position (i, j). Since X's are captured by a head camera, X's contain inference background besides eyes. Hence, Yi,j's reside in an open set that may contain unknown classes.

3.1 Eye Detection

In this step, we aim to exclude interference regions from the image, which typically contains uncontrolled noisy backgrounds. To achieve this, we use object detection techniques to localize eye sockets. By accurately locating the eye socket, we can effectively identify the crucial subregions for eye segmentation. This localization step allows us to narrow down the search space and achieve more efficient and accurate eye segmentation.

Specifically, we first train an eye detection model d(·) that takes X as input and returns a rectangular bounding box (l, t, h,w) that envelops a subregion x which only contains an eye (see the supplement for the specific training procedure). Using the bounding box, the corresponding segmentation label Y is also cropped into a subregion label y. All the resulting x and y are resized into H × W. Note that yi,j's now reside in a closed-set {0, 1, 2, 3} that corresponds to background, eye, iris, and pupil, respectively. We denote the distribution of the resulting data pair (x, y) as D.

3.2 Deterministic Eye Segmentation

We then train a deterministic segmentation network to learn a mapping fϕ : x ∈ RH×W ↦→ z ∈ RH×W×D , followed by a linear transform and softmax activation on z at each location (i, j):

where W is a 4-by-D matrix (D is the dimensionality of z) and pi,j = [p(0) i,j , p(1) i,j , p(2) i,j , p(3) i,j ]T is a probability vector, of which each component, p(c) i,j , suggests the probability that the corresponding pixel (i, j) belongs to the class c (for c = 0, 1, 2, 3). The deterministic segmentation network is trained using the following cross-entropy loss:

During inference, the deterministic segmentation model outputs a segmentation map ˆ y ∈ RH×W , of which the value at each position (i, j) is determined by

3.3 Projection Head for Uncertainty-Awareness

The deterministic segmentation model could not handle imagery blurriness and train-test domain gaps, which we find quite ubiquitous in the context of AR/VR products (see Fig. 3 for the examples). In light of this, we build an uncertaintyaware projection header on top of the backbone of the pretrained deterministic segmentation model. This header takes the feature maps of each stage of the segmentation network as input and predicts the covariance of each pixel, thus providing an estimate of the uncertainty associated with the model's predictions. The final output is the uncertainty of an entire image calculated from the pixel-level covariance.

To achieve this, we resort to probabilistic machinery by transforming the deterministic eye segmentation into a probabilistic one. Specifically, we assume that the latent code of a given pixel is no longer a deterministic one but a random variable that follows a Gaussian distribution: pθ(zi,j|x) = N(zi,j; fϕ(x, i, j), Λθ(x, i, j)), where Λθ(·) is the uncertainty-aware projection head that takes x and an image position (i, j) as input and outputs a diagonal matrix Λxi,j = Λθ(x, i, j), which captures uncertainty in the latent space. Note that this projection head can be instantiated by any network with learnable parameters θ.

To learn this posterior distribution pθ(zi,j|x, i, j), we notice that the matrix W pretrained in Eq. (1) contains prior knowledge to be utilized: the four rows of W are ideal class template vectors for background, eye, iris, and pupil, respectively. This induces a desired prior q to which pθ should be regularised. Specifically, the prior q is chosen to be the Dirac delta function, that is, q(zi,j|yi,j) = δ(zi,j - cyi,j ), where cyi,j := W T eyi,j . Here the vector eyi,j is a one-hot vector which contains one at the index yi,j and zeros at the rest of indices. Intuitively, the matrix-vector multiplication W T eyi,j takes the yi,j th row of W and transposes it as a column vector.

To perform regularization, we minimize the cross entropy (CE) between the chosen prior q and the posterior pθ at all image positions (i = 1, ..., H, j = 1, ..., W):

where each summand Li,j =

Here, Eq. (5) is obtained by plugging q(zi,j|yi,j) = δ(zi,j - cyi,j ) and pθ(zi,j|x) = N(zi,j ; fϕ(x, i, j), Λθ(x, i, j)) into Eq. (4).

After optimization, we obtain the optimal θ∗ and further, given a test ˜x, its optimal Λ∗˜xi,j that captures the dimensionwise latent uncertainty at the pixel location (i, j) of ˜x. However, we may expect one summary scalar that quantifies the uncertainty at (i, j). Next, we theoretically show that the desired scalar happens to be the trace of the optimal Λ∗xi,j .

Theorem 1. Given an image x and a position (i, j) of interest, the trace of the optimal Λ∗xi,j that minimizes each summand of the proposed loss function Eq. (4) is equal to the squared Euclidean distance between the latent code zi,j and the corresponding class center cyi,j :

where

Theorem 1 suggests that the closer the latent code zi,j gets to cyi,j , the lower the value of tr(Λ∗xi,j ) becomes. The Euclidean distance can be interpreted as the distance of the prediction at position (i, j) towards its corresponding ideal position in the latent space, and therefore the quantity tr(Λ∗xi,j )

Algorithm 1: Training Algorithm

Input: The training pairs (X, Y ) ∼ D; The pretrained detection model d(·).

Output: The deterministic segmentation network fϕ, the matrix W, the projection head Λθ

for (X, Y ) ∼ D do

(l, t, h, w) ← d(X);

x ← crop X using (l, t, h, w);

y ← crop Y using (l, t, h, w);

{ϕ, W} ← Train the eye segmentation network by minimizing Lseg in Eq. (2);

{θ} ← Train the projection head by minimizing the surrogate loss in Eq. (9);

return fϕ, W, Λθ

can be seen as the uncertainty for the prediction. In the inference stage, the category of a pixel cannot be obtained in advance, so predicting the category center to which the pixel at (i, j) belongs is an ill-conditioned problem. Our method circumvents this difficulty by estimating an uncertainty tr(Λ∗xi,j ) that mathematically measures how close a pixel under test is to its unknown class center.

3.4 Optimization Analysis

Solving the original loss function Eq. (4) requires minimizing each summand of Eq. (4) simultaneously. In practice, we find it often converges to suboptimal minima. In this section, we provide a detailed analysis and propose a surrogate loss that can circumvent this optimization difficulty by virtue of Theorem 1.

From Eq. (4), we can find the first derivative of the loss function Li,j :

There are two multiplicative factors that contribute to decreasing the loss derivative. The first factor leads to a trivial solution where the entries of Λxi,j are extremely large. In practice, we observe that the first factor dominates in finding the minimum of Eq. (4) due to random initialization of network parameters in the optimization process. This means that the entire optimization process will probably face optimization issues such as gradient vanishing. In Fig. 2 (left) we showcase the function landscape of the loss Eq. (4) as a function of Σxi,j in the two-dimensional case.

To circumvent the optimization difficulty of loss Eq. (4), we derive a surrogate loss based on Theorem 1. According to Theorem 1, the diagonals are all that we need for uncertainty estimation. Therefore, we only have to have the uncertainty module Λθ(·) to output the diagonals and make them approximate the desired quantity that is only available during training, (cyi,j -zi,j) ⊙ (cyi,j -zi,j):

where each summand ˜ Li,j is

Note that the function landscape of the surrogate loss Eq. (9) is flat and smooth as compared to the original loss Eq. (4), as shown in Fig. 2. Therefore, the optimization objective is more amenable, as will be shown in our experiments. See Algorithm 1 for the detailed training procedure.

3.5 Inference

After training, we obtain the resulting eye segmentation model fϕ, the pretrained matrix W and the projection head Λθ. Then, given a raw testing image X, we first use the pretrained eye detection model d(·) to obtain the patch x, then use fϕ and W to obtain the segmentation result ˆy. Further, besides the segmentation result, we can obtain the predicted pixel-wise covariance Λxi,j = Λθ(x, i, j) at every position (i, j). However, we expect an associated uncertainty score sunc that can summarize the uncertainty of the whole predicted result ˆy. In this section, we show how to derive this quantity from the predicted pixel-wise covariance.

Algorithm 2: Test Algorithm

Input: The testing image X; The pretrained eye model fϕ, the pretrained matrixW, the

detection model d(·); The eye segmentation projection head Σθ. Output: The predicted segmentation map ˆ y and its associated uncertainty score sunc (l, t, h, w) ← d(X); x ← crop X using (l, t, h, w); z ← fϕ(x); for (i, j) do pi,j ← softmax(Wzi,j ); ˆ yi,j ← arg maxc p(c) i,j ; Λxi,j ← Λθ(x, i, j) sunc ← ∑i,j lndet(Λxi,j ) return ˆ y, sunc We propose to use the negative logarithm of the posterior density p(z|x) as the uncertainty score sunc. The intuition is that the higher the density p(z|x) is, the lower the uncertainty score. Specifically, the posterior density can be factorized into:

Let v = (cyi,j - zi,j). Then, for the quadratic form

Substituting Eq. (17) into Eq. (11). we can derive the uncertainty estimate for the whole result from pixel-level estimate:

Therefore, our proposed uncertainty score sunc for the whole segmented result is given by

The uncertainty score sunc quantifies the predictive confidence of the model in its segmentation output by mapping the posterior probability to the real number domain R. See Algorithm 2 for the procedure of test algorithm.

4 Experiments

4.1 Datasets

There are several publicly available eye segmentation datasets [Fuhl et al., 2015; Fuhl et al., 2016b; Garbin et al., 2020;

Fuhl et al., 2021; Tonsen et al., 2016; Fuhl et al., 2016a; Fuhl et al., 2017] that have been established to aid research and algorithm development. We conduct extensive experiments to evaluate our proposed EyeSeg on multiple widely used datasets, including OpenEDS [Garbin et al., 2019], LPW [Tonsen et al., 2015], Dikablis [Fuhl et al., 2022]. These datasets provide a diverse range of real-world scenarios and variations in imaging conditions. Furthermore, we also curate a meticulously annotated dataset, Else [Fuhl et al., 2016b] to assess the performance of competing methods under more challenging conditions. The annotated labels will be open-sourced. All datasets include four classes: background, eye, iris, and pupil.

4.2 Implementation Details

To increase the diversity of training data, we use conventional data augmentation techniques such as random rotation, translation, scaling, and horizontal flipping. All experimental settings are kept the same for fair comparison. We apply gamma correction to the image during the preprocessing stage.

In the preprocessing step, to localize and identify the regions of eyes in images, we employ an object detection of which its network architecture is similar to YOLOv3 [Redmon and Farhadi, 2018]. Specifically, we employ a pruned, scaled-down version to reduce computational complexity. The network architecture is simplified into a shallower model, with the total computation reduced to 43.13M FLOPs. After this step, all regions of eyes are transformed into 96×96 images to serve as input to the segmentation model.

We utilize the lightweight segmentation network architectures DeepVOG [Yiu et al., 2019] and DenseElNet [Kothari et al., 2021] as the backbones for segmentation, which we refer to as Ours(D) and Ours(E), respectively. The input size of the network is standardized to 96 × 96, and the output category of network segmentation is set to a closed set including 4 categories (Pupil, Iris, Eye and Background). Due to eye patch extraction, our framework has a much smaller computational overhead than that of using full-images as input under the same segmentation network architecture.

Our uncertainty-aware projection header is implemented using a bottleneck architecture. Specifically, it consists of one downsampling layer and two upsampling layers with skip connections that link corresponding outputs at each layer. By incorporating these skip connections, the projection header can leverage the rich feature representations extracted by the backbone network. The downsampling layer reduces the spatial dimensions of the input, capturing high-level contextual information. The subsequent upsampling layers recover the spatial resolution, enabling fine-grained uncertainty estimation at each pixel. This architecture design ensures that our uncertainty-aware projection header can effectively capture both global and local information from the backbone network, which contributes to the robustness and reliability of uncertainty estimation.

4.3 Baseline Methods

In experiments, we consider two experimental settings: segmentation and uncertainty estimation. In segmentation, our proposed model is compared with existing top-performing and lightweight eye segmentation methods for AR/VR. We consider ad-hoc eye segmentation methods, including DeepVOG [Yiu et al., 2019], RITNet [Chaudhary et al., 2019], Pisto [Fuhl et al., 2024] and Pylids [Biswas and Lescroart, 2023], with settings consistent with the original training settings in their respective papers. Pylids [Biswas and Lescroart, 2023] is a model for keypoint detection and ellipse fitting. Pisto [Fuhl et al., 2024] executes the executable program provided in the paper. For fair comparison, only results from test datasets not used in training are evaluated. In addition, we also consider a generic segmentation method DeepLabv3+ [Chen et al., 2018] using MobileNet backbone. Notably, it has significantly higher computational complexity than the above methods.

In uncertainty estimation, baseline methods include Ensemble [Czolbe et al., 2021], EvPSNet [Sirohi et al., 2023], Dudes [Landgraf et al., 2024] and BiTrans [Wei et al., 2021]. To ensure a fair comparison, we use DeepVOG as the backbone for all the methods.

4.4 Results

In this section, we demonstrate EyeSeg's performance in terms of two experimental settings: segmentation and uncertainty estimation. Performance is measured using metrics including MIoU (Mean Intersection over Union), E1, F1 and ACC. See complete results in Appendix.

Segmentation. Tab. 1 shows the segmentation results across four datasets. On the OpenEDS dataset, our approach outperforms the state-of-the-art. This improvement is attributed to our method's ability to exclude irrelevant regions, leading to a more efficient and accurate model. These findings suggest the advantages of EyeSeg in eye semantic segmentation.

Additionally, we also evaluate the performance of our method on Else, Dikablis and LPW. It demonstrates that our proposed framework achieves better or promising per-

formance using only a segmentation network with a small amount of computation. In Fig. 4, we present a comparison of the segmentation results from our approach against RITNet and DeepVOG on the Else dataset. As shown in the figure, our method demonstrates higher accuracy and robustness across various challenging scenarios. Besides MIoU, we observe that EyeSeg also exhibits superior performance on E1, F1 and ACC in Appendix.

VOG on the Else dataset. The boxes represent the errors in segmen-

tation outcomes. It serves as a compelling testament to the outstand-

ing precision and resilience exhibited by our approach

Uncertainty Estimation. We conduct experiments to evaluate the performance of our uncertainty-aware approach on hard samples in four settings: in-domain, cross-domain, Occlusion and Blur. In the in-domain settings, we evaluate the effects of factors such as image blurriness and noise on model performance. In the cross-domain settings, we evaluate the effects of domain gaps on model performance. In the occlusion setting, we simulate real-world scenarios where parts of the eye are partially covered and assess how well the model can still segment the eye accurately. For the blur setting, we introduce varying levels of motion blur to the test samples to gauge the model's robustness in maintaining segmentation quality under degraded image conditions.

We collect all images from the dataset and rank them according to uncertainty scores generated by different models. With images of large uncertainty scores filtered out, we expect the rest of images to attain good segmentation results. To ensure fairness, all methods are evaluated based on MIoU. As demonstrated in Tab. 2, the images selected by our method show better results than those by other uncertainty-aware methods. This indicates that our proposed uncertainty-aware projection header is more effective for eye segmentation.

In Fig. 4, the uncertainty scores obtained by our method align with human cognition: 1) Our proposed uncertainty scores across cross-domain datasets positively correlate with their similarity to the training dataset Else; 2) Within each dataset, the uncertainty scores generally correlate positively with eye semantics completeness.

4.5 Ablation Study

Study I. We conduct experimental comparisons of the two optimization objectives proposed Eq. (9) in this paper. From the results presented in the Tab. 3, it can be observed that both the optimization objectives yield improvements. The results obtained from each stage of Tab. 3 demonstrate that the surrogate loss function provides a more accurate estimation of the uncertainty associated with the segmentation results.

Study II. We conduct a comprehensive comparison between our method's uncertainty estimation and segmentation performance on both the entire image and sub-regions. As shown in Tab. 3, our method consistently outperforms the other regardless of the proportion of data removed according to the uncertainty scores estimated by our framework.

5 Conclusion

We have proposed an uncertainty-aware eye segmentation framework, EyeSeg, designed to address major challenge in AR/VR: motion blur, eyelid occlusion and train-test domain gaps. We first extract eye patches and then perform uncertainty estimation and segmentation on them. The extracted eye patches satisfy the closed-set condition, reducing the model's difficulty in processing redundant background information and enhancing its robustness for cross-domain data segmentation. The patch extraction allows us to perform an uncertainty estimation Bayesian approach under which a posterior approximates a designed prior induced by the closed set. In addition, we introduce a novel theoretical-grounding approach for uncertainty estimation in this task. Our model surpasses the state-of-the-art methods.

Acknowledgments

This work is supported by the National Natural Science Foundation of China (62302167, U23A20343, 62472282, 72192821), Young Elite Scientists Sponsorship Program by CAST (YESS20240780),Shanghai Sailing Program (23YF1410500), Chenguang Program of Shanghai Education Development Foundation and Shanghai Municipal Education Commission (23CGA34).

Contribution Statement

This work was a collaborative effort by all contributing authors. Zhengyuan Peng, Jianqing Xu made equal contributions to this study and are designated as co-first authors. Jianqing Xu and Shen Li acted as project leads, offering initial ideas and guidance throughout the research process. Xin Tan, serving as the corresponding authors, are responsible for all communications related to this manuscript.