1 Introduction↩︎

Since graphene was successfully isolated [1], two-dimensional materials have rapidly become a major research focus in condensed-matter and materials physics. Experimentally, mechanical exfoliation [2] is widely used to prepare two-dimensional material samples due to its simplicity and high crystalline quality. Taking advantage of thin-film interference [3], exfoliated flakes exhibit obvious reflectance contrast and clear morphological cues under visible light, which provides a physically observable basis for pixel-level flake segmentation from microscopic images.

In recent years, semantic segmentation has been successfully applied to automated identification of two-dimensional material flakes [4], [5]. When training images are sufficiently annotated, models can learn optical contrast, edge gradients, and texture patterns of target regions, enabling precise segmentation at pixel-level. However, in real laboratory settings, images are often large, have complex backgrounds, and contain diverse impurities. Pixel-level annotation typically requires domain experts investing substantial time and effort. Consequently, the significantly increasing annotation and maintenance cost directly limits training efficiency and scalability. Purely dense-annotation-based fully supervised solutions are therefore difficult to deploy at scale in routine experimental workflows.

Weakly supervised learning offers an alternative to reduce annotation cost. As shown in 1, scribble lies between point and full mask, allowing annotators to draw coarse outlines rather than precise contours. With consistency learning or perturbation strategies, models can exploit information from unlabeled regions via sparse scribbles to recover full segmentation capability. Although scribble cost slightly more than points and have only limited information gain compared to weaker supervision forms like box, they naturally encode shape priors and often achieve a favorable trade-off between performance and annotation efficiency. It should be noted that existing scribble-supervised segmentation studies mainly focus on natural scenes [6] and medical images [7], with no systematic studies reported on two-dimensional material microscopic data.

Existing scribble-supervised segmentation methods can be roughly categorized into regularization learning [8], [9], pseudo-label learning [10]–[13], and consistency learning [14]. However, these approaches all exhibit shortcomings in the context of two-dimensional thin-section material microscopic images. Regularization learning propagates supervision signals by constraining label consistency between neighboring pixels. However, in two-dimensional material segmentation scenarios, its emphasis on low-level material features — such as color and grayscale — is susceptible to background texture interference, making it difficult to effectively integrate with high-level semantic boundary judgments. Pseudo-label learning attempts to densify supervision via self-training, yet when flake and background colors are similar and the thickness or interference color varies, the pseudo-labels generated by the model near semantic boundaries can be systematically biased and errors may accumulate over iterations. Consistency learning aligns predictions under different perturbations to regularize the model, but it mainly enforces output stability rather than directly improving boundary discriminability. At the core lies the spatial nature of scribbles: annotators typically mark high-confidence interior regions, leaving boundary pixels almost unsupervised. Boundaries are precisely the hardest part to learn because of geometric sparsity (minuscule proportion relative to interior areas), feature ambiguity (resemblance to substrate textures, contaminants, or crease shadows), and imaging degradation (uneven lighting, low-contrast transition zones). Constrained by the aforementioned bottlenecks, existing methods often compromise between boundary accuracy and regional consistency, making them ill-suited for the task of segmenting two-dimensional material sheets.

Therefore, this paper proposes TACoS (Tree-based Asymmetric Contrast Segmentation), the first annotation segmentation framework specifically designed for two-dimensional materials. It unifies weak-strong consistency learning, structured regularization, and online pseudo-label semantic guidance within a single end-to-end training framework. This approach enables efficient learning in unlabeled regions at low annotation cost while significantly enhancing boundary and structural stability. TACoS consists of three tightly coupled key designs: Unlabeled Weak–Strong Distribution alignment (UWSD), Tree Energy Regularization (TER), and Asymmetric Regional Contrastive Learning (ARCL). Specifically, UWSD enforces consistency constraints on prediction distributions between weak/strong augmentation views only on unlabeled pixels, providing dense training signals for the dominant unlabeled regions. TER builds a minimum spanning tree on backbone features, and uses tree filtering to generate a structure-aware soft reference, imposing structured regularization on unlabeled pixels to suppress foreground fragmentation and connected-component drift. ARCL completes unlabeled pseudo labels using high-confidence predictions from the weak augmentation branch, combines them with scribble annotations to form augmented labels, and applies contrastive constraints on strong-augmentation-branch representations. By turning boundary learning into a contrastive learning problem, ARCL enhances boundary-discriminative representations without additional boundary labels, compensating for the lack of direct class-level guidance at boundaries under sparse supervision.

Experimental results indicate that even with less than 0.6% annotated data, this method achieves over 96% of full-supervision performance while exhibiting superior regional coherence and boundary stability in visualization results.

In summary, the main contributions of this paper are as follows:

• For the first time, we introduce sparse supervision into the segmentation of two-dimensional materials, proposing a unified framework called TACoS that integrates semi-supervised consistency learning with structured tree energy constraints. This addresses the lack of semantic guidance and boundary degradation in unlabeled regions of thin two-dimensional material samples under sparse supervision.

• Design an asymmetric region contrast learning module to construct class prototypes in the representation space and perform region-level contrast learning. Simultaneously, focus contrast constraints on challenging pixels in unlabeled boundary regions to enhance intra-class cohesion and inter-class separation at the representation level, thereby reducing class confusion in scenarios with weak contrast edges and complex backgrounds.

• Through quantitative metrics and visual analysis, TACoS has been thoroughly validated to demonstrate superior structural consistency and boundary stability in scenarios involving weak contrast edges and complex backgrounds.

The remainder of this paper is organized as follows. 2 reviews related work on automatic segmentation of two-dimensional material micrographs, semi-supervised segmentation, and contrastive learning segmentation, outlining the relationship between our method and existing approaches along with improvements. 3 details the overall framework, training objectives, and implementation specifics of TACoS under a unified network setting. 4 reports experimental results and analysis. To ensure representativeness and reproducibility of comparisons, we adopt DINOv2 [15] as the unified backbone network, DPT [16] as the segmentation head, and align training budgets. We select typical scribble-supervised methods with different mechanisms as comparison baselines, including A\(^2\)GNN [13], URSS [8], CC4S [9], AGMM [11], and SASFormer [14]. We also include FlashInternImage [17] with a DCNv4 backbone as a strong reference point for comparative training, and conduct ablation studies to evaluate the individual contributions and coupling gains of UWSD, TER, and ARCL. Finally, 5 concludes the paper and discusses future directions.

2 Related Work↩︎

2.1 Optical microscopic Images of Two-Dimensional Materials and Automated Segmentation↩︎

Deep learning has been widely used for automated segmentation of two-dimensional material flakes, achieving progress in instance segmentation [4], [18] and semantic segmentation [19]–[21]. For instance segmentation, Masubuchi [18] built a dataset of over 2,100 optical microscopic images covering BN, graphene, MoS\(_2\), and WTe\(_2\), and integrated Mask R-CNN [22] into an electric stage microscope system to enable an end-to-end workflow for automated scanning and flake instance segmentation. For semantic segmentation, Zhang [21] systematically benchmarked multiple semantic segmentation networks on public data and proposed 2DU2-Net, providing a representative baseline for model selection and a evaluation insight for low-contrast and fragmented morphologies. Overall, public datasets still suffer from limited scale and diversity, while pixel-wise mask annotation is expensive and expert-dependent, limiting generalization and deployment.

To mitigate this, Yan [23] constructed a large-scale dataset with 7,454 images and about 30,000 annotated flake regions across multiple substrates and imaging conditions, and proposed a batch-wise self-evolving annotation ecosystem that integrates active learning, semi-supervised learning, and foundation models to reduce manual cost, establishing a stronger basis for generalizable training and evaluation.

2.2 Scribble-Supervised Semantic Segmentation↩︎

Scribble-supervised semantic segmentation aims to express the shape features of target regions with low annotation cost, thereby accomplishing segmentation tasks more efficiently. Since ScribbleSup [6] first extended point-level annotations to scribbles, related methods have continuously evolved. It is worth noting that many representative works were proposed under broader sparse-annotation settings (e.g., point annotations, bounding boxes, or mixed weak annotations), but their core mechanisms are also transferable to scribble supervision.

A\(^2\)GNN [13] propagates pseudo labels on a relational graph and suppresses weak-supervision bias using soft-edge weighting and attention propagation. Although this method can leverage edge priors to suppress errors, the effectiveness of its soft edge weights drops sharply under low contrast conditions, making it difficult for pseudo-labels to precisely align with contours. URSS [8] combines uncertainty constraints with random-walk diffusion, excluding pseudo boundaries via entropy minimization to enhance region consistency. CC4S [9] further improves pseudo label quality by introducing color regularization on this foundation. While these methods perform well in natural scenes, in two-dimensional material images where the optical features of the foreground and background are similar, random walks are highly prone to crossing semantic boundaries, resulting in pseudo-label bias. TEL [10] uses a minimum spanning tree to model low-level color and high-level semantic similarity. However, the extremely low discriminative power of low-level features interferes with the topological construction of its tree structure, compromising edge integrity. AGMM [11] and AGMM++ [12] employ class-conditional Gaussian mixture models and online expectation-maximization processes to precisely model and refine feature distributions. However, due to partial overlap between foreground and background features in the feature space, they struggle to accurately separate features on either side of the boundary. SASFormer [14] utilizes hierarchical attention to propagate labeled region information and employs an affinity consistency loss to constrain predictions. However, in the absence of significant visual cues, attention tends to diverge, and boundary pixels lack direct supervision, leading to the accumulation of mislabeling bias. Although the aforementioned methods have achieved significant results in propagation denoising and online soft supervision generation, under low contrast and complex substrates in two-dimensional-material microscopic images they can still suffer from accumulating pseudo-label bias, and boundary pixels lack direct supervision, making contour quality hard to improve reliably.

In contrast, the TACoS method proposed in this paper provides dense unlabeled training signals through weak-strong alignment, introduces a regularization strategy that includes tree-structure constraints, enhances structural stability, and prevents the propagation of errors from false labels. Furthermore, a region-level contrastive learning module is introduced to increase the distance between the specimen and the substrate background as well as other impurities in the feature space. This effectively addresses the issue of indistinguishable boundaries caused by low contrast, thereby improving the quality of the specimen outline without increasing annotation costs.

2.3 Semantic Segmentation with Semi-Supervised and Contrastive Learning↩︎

Semi-supervised semantic segmentation trains with a small amount of labeled data alongside a large volume of unlabeled data. The key is to construct reliable supervision for unlabeled pixels and prevent bias amplification. Classic approaches like CCT [24] enhance training stability by constraining unlabeled regions through cross-consistency constraints across multi-branch and multi-perturbation predictions. CPS [25] reinforces consistency and extends effective training signals by employing dual networks mutually supervised with pseudo labels. In addition, the recently proposed B\(^3\)CT [26] further designs a three-branch learning architecture that achieves domain-robust semantic segmentation by mining unlabeled target signals, providing a new paradigm for leveraging unlabeled information under complex feature distributions. Meanwhile, contrastive learning enhances intra-class compactness and inter-class separation in the representation space. For instance, CIPC [27] introduces global context through cross-image pixel-level comparisons, while STEGO [28] distills semantic clusters under fully unlabeled conditions by leveraging pre-trained self-supervised features to capture consistent structural relationships within and across images. Combined with contrastive loss and structured graph optimization, it achieves consistent and detailed unsupervised semantic segmentation results.

Although this paper focuses on scribble-based weak supervision, the utilization of a large number of unlabeled pixels aligns closely with the semi-supervised setting. we propose the single-stage framework TACoS, comprising three modules, each introducing its reference source and design rationale: UWSD is inspired by the weak augmentation pseudo-label and strong-augmentation consistency alignment of the classic semi-supervised model FixMatch [29] and leverages data-augmentation strategies from UniMatch V2 [30] to better utilize the large proportion of unlabeled pixels. TER is motivated by the minimum spanning tree modeling in TEL [10], using tree structures to characterize pixel affinity relationships and a structure-aware soft reference to suppress fragmentation and drift. Inspired by contrastive learning model, ARCL draws from the region-prototype and hard-sampling mechanism of ReCo [31] and the boundary pixel pair contrast with confidence-guided prediction in GCBL [32]: it completes supervision with high-confidence predictions from the weak augmentation branch and drives regional contrastive learning, focusing representation constraints on uncertain and boundary neighborhoods to improve boundary discriminability.

3 Method↩︎

3.1 Task Setup↩︎

This paper targets semantic segmentation of two-dimensional material flakes in optical microscopic images, adopting a default binary classification setup, the foreground representing flake regions, while the background encompasses blank substrates, multilayer non-target areas, impurities, and other elements. Let the input image be

\[x\in\operatorname{R}^{H\times W\times C}.\]

Segment network output pixel-level category probability maps

\[p\in[0,1]^{H\times W\times K},\;\;K=2.\]

In our dataset, dense masks take values in 0, 1; for the scribble setting, the label mask uses 0, 1, 255, where 255 denotes unknown pixels. We partition the pixel set \(\Omega\) into a labeled pixel set \(\Omega_L\) and an unlabeled pixel set \(\Omega_U\), such that \(\Omega_L \cup \Omega_U = \Omega\) and \(\Omega_L \cap \Omega_U = \emptyset\). This explicit modeling of pixel-level labeling deficiencies serves as the core starting point for the subsequent design of TACoS.

3.2 TACoS: Unified Scribble Segmentation Framework↩︎

As shown in 2, the proposed TACoS framework adopts a dual-branch joint optimization paradigm to recover complete segmentation masks for two-dimensional materials from extremely sparse scribble annotations. Given an input image, the network feeds its weak and strong augmentation views into an encoder-decoder with shared weights. To fully leverage information from unlabeled regions, TACoS tightly couples three core modules during end-to-end training: The Unlabeled Weak-Strong Distribution Alignment (UWSD) module forces the strong enhancement branch’s predictions to align outputs from the weak branch, providing dense pixel-level consistency supervision signals for unlabeled areas. The Tree Energy Regularization (TER) module combines low-level details with high-level semantic features to construct minimum spanning trees and perform two-stage tree filtering, generating soft references with global awareness that impose strict spatial structural regularization constraints on unlabeled pixels. The Asymmetric Region Contrast Learning (ARCL) module fuses online pseudo labels with sparse scribbles and enforces region-level asymmetric contrast constraints in the feature space of the strong enhancement branch, significantly improving boundary discrimination accuracy. Ultimately, these three constraints are jointly optimized with standard scribble loss within a unified objective function, achieving a stable transition from sparse supervision to high-quality dense prediction. Unlike multi-stage approaches, TACoS employs a strict single-stage design. This means that all optimization objectives are jointly minimized in a single continuous training process, completely eliminating the need for offline pseudo-label generation, iterative model retraining, and separate fine-tuning stages.

3.2.1 UWSD: Unlabeled Weak–Strong Distribution Alignment↩︎

In weakly supervised semantic segmentation, the key to effectively utilizing unlabeled pixels lies in constructing a robust consistency regularization framework. Drawing inspiration from UniMatch V2’s design, this paper adopts a weak-strong dual-branch architecture: leveraging the model’s high-quality predictions on images with simple transformations (weak augmentations) to guide its learning of representations under complex transformations (strong augmentations) of the same images. This training approach not only significantly enhances the model’s robustness against noise and perturbations but also challenges it to discover invariant feature representations across different image perspectives.

To extend consistency regularization from the image level to the pixel level, given a microscopic image \(x\), we construct the weak augmentation view and the strong augmentation view within the same cropped field of view during training:

\[x^w=T_{geo}(x),\;\;x^s=T_{app}(T_{geo}(x)),\] here, \(T_{geo}(\cdot)\) denotes geometric transformations, including random scaling, random cropping to fixed dimensions, and random horizontal flipping. Input \(x\) undergoes this transformation to generate the weak augmentation branch input \(x^w\), with the scribble annotation label \(y\) transformed synchronously. Subsequently, the strong augmentation branch further applies appearance perturbations \(T_{app}(\cdot)\) relevant to microscopic imaging on top of this geometric transformation: it performs color jittering (brightness / contrast / saturation / hue) with a certain probability, and applies grayscaling and random blurring with a low probability. This simulates imaging variations such as fluctuations in light intensity, color temperature, and slight defocusing. Thus, the input \(x\) undergoes both geometric transformation and appearance perturbation to generate the input \(x^s\) for the strong augmentation branch. The geometrically transformed label \(y\) serves as the direct supervision signal for this branch, while \(x^w\) and \(x^s\) are utilized for consistency regularization learning.

To ensure consistency regularization acts solely on the unlabeled pixel set \(\Omega_U\). In unlabeled regions, UWSD employs the prediction distribution of the weak augmentation view as a soft objective, constraining the predictions of the strong augmentation view to align with it. This suppresses the propagation of apparent perturbation bias from strong augmentation onto unlabeled pixels.

Let \(d(p,q) = 1 - \cos(p,q)\), then we have

\[d(p,q) = 1-\frac{\langle p,q\rangle}{\|p\|_2\,\|q\|_2+\epsilon}.\]

We use the weak-augmentation prediction as the alignment target and stop its gradients from backpropagating.

\[\mathcal{L}_{\text{UWSD}}=\frac{1}{|\Omega_U|}\sum_{i\in\Omega_U}d\!\left(\mathrm{sg}\!\left(p_w(i)\right),\,p_s(i)\right).\]

Stopping gradient backpropagation is necessary because the weak-augmentation branch only provides a relatively stable soft target and is not updated by this loss. This prevents the weak and strong augmentation branches from chasing each other, which could otherwise cause distribution drift or training collapse. This is particularly important under sparse supervision, since \(\Omega_U\) typically accounts for the vast majority of pixels in an image.

The weak/strong augmentation views share the same set of network parameters and utilize a common encoder and decoder, with their branching structure following the overall paradigm of UniMatch V2.

3.2.2 TER: Tree-Energy Regularization↩︎

UWSD can provide dense consistency signals on \(\Omega_U\), but it remains fundamentally an alignment at the distribution level without explicitly constraining structural relationships between pixels. To address this, we introduce TER, which constructs a minimum spanning tree (MST) in feature space to capture pixel affinity relationships and generates structure-aware online soft references based on this. To address the optical properties of two-dimensional material flakes, we extract shallow and deep feature representations from DINOv2 to construct a dual-tree architecture. The shallow features capture microscopic textures and color gradients to anchor boundaries, while the deep features encode high-level semantics to suppress background noise. This dual-tree mechanism ensures that the generated soft references possess both microscopic geometric accuracy and macroscopic semantic consistency.

Let the low-level (Layer 2) and high-level (Layer 11) features extracted from the strong-augmentation branch be

\[f_{low}\in\operatorname{R}^{H_l\times W_l\times D_l},\;\;f_{high}\in\operatorname{R}^{H_h\times W_h\times D_h}.\]

Based on the aligned features, we define a pixel graph where all pixels form the vertex set and the feature distances between adjacent pixels form the edge set. By successively removing edges with the largest weights while preserving the graph’s connectivity, we can construct a minimum spanning tree:

\[\mathcal{G}=\operatorname{MST}\left(\Omega,\mathrm w_{ij}=\parallel f(i)-f(j)\parallel_2\right),\] where \(f(\cdot)\) takes the values \(f_{low}\) and \(f_{high}\), respectively, we construct two minimum spanning trees \(\mathcal{G}_{low}\) and \(\mathcal{G}_{high}\) in the low-level and high-level feature spaces, respectively.

To capture long-range dependencies between pixels, we map tree distances to an affinity matrix:

\[A_{ij}=\exp\left(-\frac{D_{ij}}{\sigma_{aff}}\right),\] where \(D_{ij}\) denotes the sum of path edge set distances between any two points in the minimum spanning tree \(\mathcal{G}\), and \(\sigma_{aff}\) represents the scale parameter.

Given a strong augmentation branch prediction \(p_s\), we apply the tree filter operator \(\mathcal{T}(\cdot)\) to the affinity matrix \(A\), yielding a structure-aware soft reference:

\[a=\mathcal{T}_{high}(\mathcal{T}_{low}\left(\Omega,\;p_s,A)\right),\;\;a(i)\in[0,1]^K,\] among these, the filtering operator \(\mathcal{T}\) represents the result of normalizing the weighted sum of the affinity \(A_{ij}\) across all image pixels \(\Omega\), weighted by the prediction probability \(p_s\). This process can be viewed as adaptive smoothing of the prediction distribution under pixel affinity constraints, ensuring that highly similar pixels remain consistent in the prediction space. Unlike offline-generated pseudo labels, \(a\) is computed online from the current feature structure and current prediction, dynamically updated at each iteration. This provides structure-aware semantic guidance for \(\Omega_U\) without increasing manual annotation costs, while effectively suppressing connected domain drift and prediction bias.

On the unlabeled pixel set \(\Omega_U\), we minimize the discrepancy between the strong-augmentation-branch prediction and the soft reference:

\[\mathcal{L}_{TER}=\frac{1}{\mid\Omega_U\mid}\sum_{i\in\Omega_U}\parallel p_s(i)-\operatorname{sg}(a(i))\parallel_1.\]

3.2.3 ARCL: Asymmetric Regional Contrast Learning↩︎

Although UWSD and TER provide dense consistency signals and structured regularization on the unlabeled pixel set \(\Omega_U\), the representational separability at the category boundary may still be insufficient for models when dealing with weakly contrasting edges and complex backgrounds. To this end, we introduce ARCL, imposing constraints on boundary neighborhoods from both the representation space and decision space. On one hand, performing region-level contrastive learning on deep decoder features. On the other, applying an exclusion penalty based on asymmetric distance sampling to boundary neighborhoods on strong augmentation branch logits. Both share the same filtering logic for valid pixels and difficult samples, thereby forming a consistent boundary enhancement mechanism.

We first construct the augmented label \(\widetilde{y}\) using the weak augmentation branch prediction \(p_w\). For pixel \(i \in \Omega\):

\[\widetilde{y}(i)=\left\{\begin{matrix}y(i),&i\in\Omega_L,\\\arg{\max}_kp_w^k(i),&i\in\Omega_U\land{\max}_kp_w^k(i)\geq\tau_w,\\\mathrm{255},&\mathrm{otherwise},\\\end{matrix}\right.,\] here, \(\tau_w\) corresponds to the weak confidence threshold. In this binary classification segmentation task, \(k\) takes the values 0 or 1. This yields the effective pixel mask.

\[m_V\left(i\right)=\mathbb{I}\left[\widetilde{y}\left(i\right)\neq\mathrm{255} \right].\]

Unlike label generation, sampling and weight computation in contrastive learning are entirely based on strong augmentation branch predictions. A stricter strong confidence threshold \(\tau_s\) is employed to filter out high-confidence samples. We focus the candidate set on the collection \(\mathcal{S}\) of challenging pixels that are classified as foreground by augmented labels but exhibit sharp confidence drops or misclassification under strong perturbations.

\[\mathcal{S}=\left\{ i \;\middle|\; m_V(i)=1 \;\land\; \max\limits_{k} p_s^{k}(i) < \tau_s \right\}.\]

For each pixel \(i \in \Omega\), let \(\mathbf{r}_i\) denote the depth-aggregated feature extracted by the segmentation network from the strong augmentation view. On \(\mathcal{S}\), the contrastive loss is defined as:

\[\mathcal{L}_{\text{Region}} =-\frac{1}{|\mathcal{S}|}\sum_{i\in\mathcal{S}}\log\frac{\exp\!\left(\langle r_i, c_{\tilde{y}(i)}\rangle/T\right)}{\sum_{k}\exp\!\left(\langle r_i,c_k\rangle/T\right)},\] here, \(c_k\) denotes the mean feature vector of valid pixels in class \(k\) within the current input batch (the class \(k\) regional prototype), while \(c_{\tilde{y}(i)}\) represents the target class prototype corresponding to pixel \(i\) based on the augmented label \(y(i)\), and \(T\) represents the temperature coefficient. This loss function promotes intra-class cohesion and expands inter-class separation at the representation level.

Meanwhile, within the decision space, ARCL directly imposes constraints on the boundary neighborhoods of strong branch logits \(z_s\). Since features in boundary neighborhoods are typically highly smooth, traditional symmetric local sampling often leads to gradient cancellation or misdirection. To address this, we adopt an asymmetric distance-based boundary search strategy.

The label map for boundary supervision is denoted as \(y_B\). In the implementation, a full mask is uniformly applied to extract boundaries. The morphological gradient based on the label map \(y_B\) yields the geometric boundary set \(\mathcal{B}\). In practice, we select boundary pixels for supervision only within the set of challenging pixels \(\mathcal{S}\), i.e., \(\mathcal{E} = \mathcal{B} \cap \mathcal{S}\). For each pixel \(i\in\mathcal{E}\) on the boundary, we extract a pair of asymmetric samples along its normal vector: one closer to the boundary (\(i^+\), the easily confused point) and one farther away (\(i^-\), the high-confidence point). If the model assigns consistent semantic categories to points on both sides of the boundary, a penalty is applied:

\[\mathcal{L}_{\text{Boundary}}=\frac{1}{|\mathcal{E}|}\sum_{i\in\mathcal{E}}\left(1+w(i)\cdot\frac{\langle z_s(i^+),\, z_s(i^-)\rangle}{\|z_s(i^+)\|_2\,\|z_s(i^-)\|_2}\right),\] the weight

\[w(i)=\exp(-\ell_{CE}(z_s(i),y_B(i))).\]

Adaptive adjustment via cross-entropy loss ensures that strong repulsive forces are applied to proximal blurred pixels only when the reference point prediction is highly accurate. If the reference point itself exhibits significant uncertainty, the penalty strength exponentially decays, thereby preventing the propagation of erroneous gradients.

Ultimately, the overall objective function for ARCL is as follows.

\[\mathcal{L}_{\text{ARCL}}=\lambda_r\,\mathcal{L}_{\text{Region}}+\lambda_b\,\mathcal{L}_{\text{Boundary}}.\]

3.2.4 Overall Loss and Optimization Objective↩︎

Combining sparse supervised loss, UWSD, TER, and ARCL, the final optimization objective is:

\[\mathcal{L}=\mathcal{L}_{sup}+\lambda_c\mathcal{L}_{UWSD}+\lambda_t\mathcal{L}_{TER}+\mathcal{L}_{ARCL},\] the \(\mathcal{L}_{sup}\) is calculated solely on \(\Omega_L\), while \(\lambda_c\), \(\lambda_t\), \(\lambda_r\), \(\lambda_b\) are used to balance the contributions of the UWSD and TER sub-terms with those of the ARCL sub-term. Additionally, all alignment targets and pseudo labels generated by weak augmentation branches halt gradient backpropagation to ensure training stability.

UWSD, TER, and ARCL are jointly optimized within the same training process, simultaneously improving regional consistency, structural stability, and boundary discriminability under sparse supervision conditions.

3.3 Network Architecture↩︎

The overall network adopts an encoder-decoder architecture, where the encoder shared by both the strong and weak augmentation branches utilizes DINOv2 as its backbone to leverage its exceptional representational capabilities acquired through large-scale data pre-training. To comprehensively cover the structural spectrum ranging from low-level textures to high-level semantics, we extract intermediate features from the \(i \in \{2, 5, 8, 11\}\)th layers among all 12 feature layers of DINOv2. This sampling method is a classic configuration of the DPT architecture. In the context of this task, shallow-layer features are sensitive to local contrast and accurately capture color shifts caused by varying layer depths and optical interference from the substrate. Deep-layer features, on the other hand, aggregate global information, abstracting the overall morphology of the material wafer and resisting interference from local surface impurities. This systematic sampling provides multi-scale representations to address the complexity of real-world two-dimensional material micrographs.

The output features of the encoder on the \(\ell\)th layer are

\[f^{\left(\ell\right)}={\rm \operatorname{Enc}}_\theta^{\left(\ell\right)}(x)\in\operatorname{R}^{H_\ell\times W_\ell\times D_\ell}.\]

The multi-level features \(\left\{ f^{\left(\ell\right)} \right\}_{\ell=1}^{L}\) undergo \(1 \times 1\) convolution projection before being fed into a DPT-based decoder for fusion. The decoder performs deep-to-shallow feature alignment and deep fusion through four levels of RefineNet modules, ultimately upscaling the features by a factor of 14 via bilinear interpolation to generate pixel-level logits:

\[z={\rm \operatorname{decoder}}_\theta({f^{\left(\ell\right)}})\in\operatorname{R}^{H\times W\times K},\;\;p=Softmax(z).\]

This yields the weak-augmentation prediction \(p_w\) and the strong-augmentation prediction \(p_s\), respectively:

\[p_w=\operatorname{F}_\theta(x^w),\;\;{p}_s=\operatorname{F}_\theta(x^s).\]

Network training simultaneously leverages four types of supervisory signals. First, it computes sparse labeling supervision loss on the labeled pixel set \(\Omega_L\). Second, it applies UWSD on the unlabeled pixel set \(\Omega_U\) to ensure weak/strong augmentation predictions align across unlabeled regions. Third, it introduces TER. This module utilizes a tree structure generated from encoder features that provides structured soft references, delivering online structural constraints and semantic guidance for unlabeled areas. Furthermore, ARCL serves as an essential framework component, enhancing boundary-case handling at both representation and decision levels. On one hand, it constructs augmented labels using weak augmentation branch predictions to expand effective training pixels. On the other hand, it performs sampling and constraints for contrastive learning based on strong augmentation branch probabilities, while combining full-mask construction with boundary supervision on strong augmentation branch logits to impose explicit constraints on boundary neighborhoods. This enhances intra-class cohesion, inter-class separation, and boundary stability.

4 Experiments↩︎

4.1 Construction and Partitioning of Sparse Annotation Datasets for Two-Dimensional Materials↩︎

To comprehensively evaluate the effectiveness and generalization capability of the proposed method in the task of segmenting microscopic images of two-dimensional materials, this study employed two representative independent datasets: the large-scale multi-scenario segmentation dataset proposed by Yan [23] and the conventional microscopic segmentation dataset proposed by Uslu [19].

The dataset proposed by Yan encompasses two materials — graphene and MoS\(_2\) — and covers multiple fields of view, magnifications, and resolutions, along with typical interfering backgrounds. It represents the largest and most diverse dataset in the field of two-dimensional material flake segmentation to date. It comprises 7,454 microscopic images, including 4,942 graphene images and 2,512 MoS\(_2\) images, with approximately 30,000 annotated flake regions. The dataset covers common resolutions such as \(1360 \times 1024\), \(3088 \times 2056\), and \(4272 \times 2848\), featuring a wide range of flake scales. The dataset proposed by Uslu was collected by an automated detection platform and comprises 2,299 microscopic images, including 1,787 graphene images and 512 WSe\(_2\) images, all uniformly resolved at 1920×1200 pixels.

Building upon the fully annotated labels provided in the original dataset, this paper constructs a two-dimensional material segmentation scribble annotation dataset using the automated scribble generation paradigm from Scribbles for All [33]. Regarding data processing and partitioning strategies, we implement reasonable adaptations tailored to the characteristics of different data sources: For the Yan dataset, we retain the original mutually exclusive partitioning scheme: graphene training (3,953 images) and validation (989 images), MoS\(_2\) training (2,007 images) and validation (505 images). The total dataset comprises 5,960 training images and 1,494 validation images, maintaining an approximate training-to-validation ratio of 4:1. For the Uslu dataset, considering the original multi-class classification task, we uniformly map it to the binary classification task defined in this paper. All materials across all layers are labeled as foreground, while the remainder are labeled as background. Additionally, recognizing the requirement for large training datasets in deep learning models, we swap the small training set and large validation set originally designed for traditional machine learning. The final division is: graphene training set (1,362 images), validation set (425 images); WSe\(_2\) training set (420 images), validation set (92 images). The total dataset comprises 1,762 training images and 517 validation images, with a training-to-validation ratio of approximately 7:2.

To visually demonstrate the extreme sparsity of scribbles, details the pixel coverage statistics across both training datasets. The data distribution and annotated pixel coverage are shown in 1. 3 illustrates examples of scribbles within the dataset. It can be observed that the total pixel coverage of scribbles consistently falls below 0.4%, fully aligning with the sparse supervision framework characterized by extremely low annotation costs proposed in this paper.

4.2 Implementation Details↩︎

Our network uses DINOv2 Small as its backbone, training without freezing the backbone parameters. The task is set as binary semantic segmentation, ignoring labels with value 255, and OHEM [34] is employed to mitigate class imbalance. For data processing, we followed the setup described in the UniMatch V2 framework [20]. Training inputs are randomly cropped to 518×518 undergo geometric augmentation via random scaling between 0.8 and 1.2, along with random horizontal flipping applied with a 0.5 probability. Weak and strong augmentation views were constructed for consistency learning. Strong augmentation branch additionally applies color and blur jitter with a probability of 0.8. The ColorJitter parameters (brightness, contrast, saturation, hue) are set to 0.2, 0.2, 0.2, and 0.1 respectively; random grayscale probability of 0.05. Gaussian blur probability of 0.2, with the blur radius \(\sigma_{\text{blur}}\) randomly sampled within the range of 0.1 to 2.0. These specific data augmentation strategies and hyperparameter ranges are designed to accurately simulate common imaging degradation phenomena in the microscopic characterization of two-dimensional materials, such as color temperature drift caused by aging of the base light source and slight defocusing caused by the movement of the motorized stage.

In the TEL module, the Gaussian kernel bandwidth parameter \(\sigma_{\text{aff}}\) used to construct the minimum spanning tree affinity matrix is set to the empirical value of 0.1. In the ARCL module, the weak confidence threshold \(\tau_w\) for pseudo-label filtering is set to a fixed value of 0.7; the strong confidence threshold \(\tau_s\) for difficult sample mining is set to a fixed value of 0.97. The temperature coefficient \(T\) of the region comparison loss is set to 0.5, with 256 queries and 512 negative keys sampled online for each class.

The optimizer employs AdamW[35] with a base learning rate of \(5 \times 10^{-6}\). Non-backbone parameter groups use a 40x learning rate multiplier. The learning rate follows a poly decay with a fixed power of 0.9. Training runs for 50 epochs with a single-GPU batch size of 2. Two NVIDIA GeForce RTX 3090 24GB GPUs are utilized, with DDP and AMP enabled. The validation phase employed sliding-window inference with a window size of \(518 \times 518\) and stride of \(448 \times 448\). mIoU served as the evaluation metric, and the model achieving the highest mIoU on the validation set was saved. In addition, the boundary IoU is evaluated simultaneously, and the contour expansion ratio is set to 2% of the image’s diagonal dimension.

4.3 Quantitative Comparison with State-of-the-art Methods↩︎

To validate the effectiveness of the proposed method, we first constructed a scribble dataset on the large-scale multi-scenario dataset introduced by Yan [23] and conducted thorough quantitative evaluations. 2 presents quantitative results on the graphene and MoS\(_2\) datasets, respectively. To ensure fairness in the baseline comparison and the reproducibility of the results, all comparison methods included in the evaluation used the same backbone network, input image dimensions, number of training iterations, optimizer hyperparameters, and validation strategy. As mentioned earlier, existing state-of-the-art scribble-supervised semantic segmentation methods can be categorized into regularized learning (URSS, CC4S), pseudo-label learning (TEL, A\(^2\)GNN, AGMM, AGMM++), and consistency learning (SASFormer) paradigms. In the context of two-dimensional material microscopic images, to mitigate issues such as noise, under-constraint, and feature confusion arising from sparse annotations, these high-performing representative works often rely on multi-stage training strategies. However, this implicit trade-off comes at the cost of higher training overhead and engineering complexity. In contrast, the proposed TACoS employs a single-stage end-to-end training approach that not only avoids complex workflows but also demonstrates superior performance over multi-stage baseline methods.

Under the graphene scribble supervision setting, the baseline method DPT achieves a mIoU of 77.86\(\%\) using the DINOv2 backbone with only supervised cross-entropy. The multi-stage training URSS improves performance to 82.35\(\%\). Using the same backbone network, our proposed TACoS achieves an mIoU of 84.09\(\%\) under single-stage training without requiring additional data, representing a 6.23\(\%\) improvement over the baseline and further surpassing URSS by 1.74\(\%\). Under the MoS\(_2\) scribble supervision setting, the DPT baseline achieves mIoU of 78.78\(\%\). Among existing methods, CC4S reaches 83.29\(\%\), while TACoS achieves the best performance at 85.20\(\%\), improving by 6.42\(\%\) over the baseline and surpassing CC4S by 1.91\(\%\).

TACoS significantly narrowed the gap with the fully supervised upper bound on both datasets, falling only 2.24\(\%\) and 3.15\(\%\) short of fully supervised DPT for graphene and MoS\(_2\), respectively. This fully demonstrates that TACoS effectively breaks the implicit trade-off between boundary accuracy and regional consistency in existing methods, maintaining robust segmentation performance and outstanding generalization capabilities even under the stringent constraints of extremely low annotation costs and single-stage training.

To further evaluate the accuracy of edge detection, this paper introduces Boundary IoU [36] as a supplementary metric. As shown in 2, TACoS achieved the best Boundary IoU performance on graphene and MoS\(_2\) at 66.95% and 62.43%, respectively, outperforming methods such as SASFormer and CC4S, thereby validating the advantages of the ARCL and TER modules in enhancing edge discrimination.

4.4 Performance evaluation on the dataset proposed by Uslu↩︎

We maintained the training settings and constructed scribble annotations on the two-dimensional material microscopic image dataset comprising graphene and WSe\(_2\) proposed by Uslu [19], followed by evaluation. We report three categories of results: DPT baseline as the lower bound, TACoS as the proposed strategy, and fully annotated training as the upper bound reference. 3 summarizes the results. Compared to the baseline, our method achieved stable gains on both datasets (2.99% improvement on graphene and 4.66% improvement on WSe\(_2\)), and significantly narrows the gap with the fully labeled training upper bound. It should be noted that due to the particularly severe imbalance in the number of pixels between foreground and background regions in this dataset, the model’s overall performance on the Uslu dataset is significantly inferior to its performance on the Yan dataset. Nevertheless, our method still demonstrates a certain degree of effectiveness in addressing such extremely imbalanced data.

6pt

Furthermore, on the Boundary IoU metric, TACoS achieved 67.88% and 68.12% on graphene and WSe\(_2\), respectively — a significant improvement over the baseline (63.74% and 61.50%) — demonstrating that the model retains excellent physical boundary preservation capabilities even under extreme class imbalance and complex background interference.

4.5 Qualitative Demonstration of TACoS and State-of-the-art Methods↩︎

To visually demonstrate the segmentation performance of the TACoS framework, 4 and 5 present qualitative comparisons between our method and existing state-of-the-art approaches on datasets proposed by Yan and Uslu, respectively. Whether on the large-scale multi-scenario dataset containing graphene and MoS\(_2\) (4) or on conventional microscopic image datasets (5), TACoS demonstrates significantly superior segmentation performance compared to the baseline model and existing state-of-the-art weakly supervised methods. Visual comparisons reveal that TACoS substantially reduces prediction errors, generating segmentation masks with clearer boundaries and high consistency with ground truth labels. This demonstrates that even under extremely sparse scribble supervision, the method exhibits robust generalization and recognition capabilities, with overall segmentation quality approaching the upper limit of fully supervised benchmarks.

Although TACoS significantly bridges the performance gap between weakly supervised and fully supervised models, extreme microscopic imaging conditions still present specific challenges. As shown in 6, misclassifications primarily occur in regions with complex background interference and low optical contrast. The visibility of two-dimensional materials is governed by thin-film interference; when the path length difference is small, the interference contrast drops sharply, blurring the physical boundary between the material edges and the substrate. In sparse labeling, the absence of dense pixel-level boundary constraints reduces the network’s sensitivity to gradient changes along these edges, leading to boundary misalignment. Furthermore, microscopic impurities, substrate defects, or dielectric layers of varying thickness — whose optical properties resemble those of the target material — often trigger false-positive predictions.

4.6 Ablation Studies↩︎

Given that the dataset proposed by Yan is the most representative in terms of data scale, field-of-view diversity, and complex background interference, all ablation experiments and hyperparameter analyses in this paper are conducted on this dataset to comprehensively validate the effectiveness and operational mechanisms of each core component within the TACoS framework.

Loss construction and effectiveness of key terms. As shown in 4, when using only the sparse supervision baseline, the mIoU for graphene and MoS\(_2\) were 77.86% and 78.78%, respectively. After introducing UWSD, performance significantly improved to 81.42% and 81.83%, corresponding to gains of +3.56% and +3.05%. This demonstrates that weak–strong consistency effectively constrains prediction stability in unlabeled regions, constituting the primary source of overall improvement. When only TER is introduced, graphene shows a modest improvement (+0.47%), while MoS\(_2\) demonstrates a more pronounced gain (+3.81%), indicating that tree energy constraints are more effective at propagating structure to unlabeled pixels in the MoS\(_2\) domain. Furthermore, simultaneously enabling UWSD and TER achieves 82.40% and 82.93%, respectively, yielding additional gains of +0.98% and +1.10% relative to UWSD alone, demonstrating their complementary nature. Introducing ARCL alone yields 81.04% and 80.19% (+3.18% and +1.41%), confirming that regional-level representation constraints enhance discriminative power. Finally, simultaneously enabling all three achieved optimal scores of 84.09% and 85.20%, representing further improvements of +1.69% and +2.27% compared to UWSD and TER alone. MoS\(_2\) showed more significant gains, indicating that ARCL provides stronger representation compensation in more complex and confusing domains.

3pt

This paper does not perform independent ablation analyses of \(\mathcal{L}_{\text{Region}}\) and \(\mathcal{L}_{\text{Boundary}}\) in the ARCL module. At the algorithmic level, the two are tightly coupled within the same dynamic hard sample pool \(\mathcal{S}\); forcibly removing either one would cause a shift in sample selection during backpropagation, so it is more reasonable to evaluate them as an indivisible module.

Comparison with Naive EMA and UniMatch v2 Strategies. Beyond the three components above, we further examine whether teacher-student style strategies yield additional gains. As shown in 5, adding either the naive EMA or UniMatch v2 strategy to the UWSD baseline yields slight improvements on graphene (+0.34% and +0.42%, respectively), but performance declines on MoS\(_2\) (0.51% and 1.00%, respectively). This phenomenon indicates that under the material domain and annotation format of this task, directly applying generic EMA or UniMatch v2 perturbation strategies is unstable. Such approaches may amplify pseudo-label noise or cause teacher lag, thereby offsetting potential gains. Therefore, this paper ultimately adopts a task-matched combination of UWSD, TER, and ARCL, rather than relying on naive EMA or directly transferred UniMatch v2 strategies.

3pt

Hyperparameter analysis. We performed grid searches on the weights of each loss term across multiple candidate values, reporting representative combinations and results in 7. To avoid evaluation bias, all hyperparameters were selected solely on the training and validation sets: the fixed network architecture, backbone network, data augmentation, and training budget remained constant. The optimal combination selected on the validation set based on the primary metric mIoU was adopted as the final configuration. Considering the significant differences in contrast distribution, boundary discernibility, and background texture interference between graphene and MoS\(_2\) in optical microscopic images, we report the weight searches and optimal settings for each material separately to reflect the method’s stable operating range across different material domains. Overall, the final adopted hyperparameters are:

• graphene: \(\lambda_c=5,\;\lambda_t=1,\;\lambda_r=0.002,\;\lambda_b=0.00025\).

• MoS\(_2\): \(\lambda_c=5,\;\lambda_t=1,\;\lambda_r=0.0025,\;\lambda_b=0.00025\).

The difference of several orders of magnitude between these loss weights acts as an inversely proportional scaling mechanism. Specifically, the global loss is computed over a massive unlabeled grid, resulting in weak single-pixel gradients; conversely, the local loss generates extremely strong, concentrated gradients along sparse, hard-to-classify regions and edges. Assigning a minimal weight to the local loss prevents these massive local gradients from overshadowing the smooth global alignment, thereby ensuring balanced optimization.

5 Conclusion↩︎

This paper presents TACoS, a novel, single-stage framework designed to address the significant bottleneck of expensive pixel-level annotation in the automated high-throughput screening of two-dimensional materials. By synergistically integrating weak-strong consistency learning, tree-energy regularization, and asymmetric regional contrastive learning, TACoS effectively propagates structural and semantic information from highly sparse scribble annotations to unlabelled regions.

The primary strength of our approach lies in its remarkable annotation efficiency and boundary precision. TACoS achieves over 96% of the performance of fully supervised methods using less than 0.6% annotated data. For the materials science community, this provides a highly scalable and practical solution. Researchers can now deploy automated flake identification systems without the prohibitive cost of dense manual labeling, thus significantly accelerating the discovery and characterization of novel two-dimensional materials and heterostructures.

Despite these advantages, our framework has limitations. In extreme scenarios characterized by exceedingly low optical contrast or highly complex, textured substrates, the model can still produce false positives or false negatives. The asymmetric regional contrastive learning (ARCL) module relies on relatively clear morphological boundaries; thus, when background impurities perfectly mimic the contrast profiles of target flakes, localized feature confusion can still occur.

Moving forward, future research will focus on extending the TACoS framework to handle multi-class instance segmentation, particularly for identifying overlapping heterostructures of different two-dimensional materials. Additionally, exploring active learning strategies to intelligently suggest the most informative scribble locations could further minimize annotation costs. Finally, extending TACoS to advanced imaging modalities such as Scanning Electron Microscope (SEM) or Atomic Force Microscope (AFM) represents a highly promising path toward a universal foundational model for material segmentation. To integrate these new domains, we propose shifting the distance metric in SEM images from optical color to high-frequency texture descriptors, and directly incorporating physical height differences in AFM data to construct three-dimensional topological constraints that are strictly aligned with actual atomic steps.

Author Contributions↩︎

Jiabei Chen: Conceptualization, Methodology, Software, Formal analysis, Investigation, Data curation, Visualization, Writing – original draft. Liping Zhang: Conceptualization, Methodology, Supervision, Writing – review & editing. Jiang-Bin Wu: Resources, Data curation. Zhongming Wei: Resources, Data curation. Enhao Ning: Writing – review & editing. Su Yan: Methodology, Software. Weijun Li: Supervision, Project administration. Ping-Heng Tan: Resources, Data curation. Xin Ning: Conceptualization, Supervision, Project administration, Funding acquisition, Writing – review & editing.

Declaration of competing interest↩︎

The authors declare that they have no known competing financial interests or personal relationships that could have appeared to influence the work reported in this paper.

Acknowledgments↩︎

This work was supported by the National Key Research and Development Program of China (Grant No. 2024YFA1409700) and the CAS Project for Young Scientists in Basic Research (Grant No. YSBR-090).

Data and Code Availability↩︎

The dataset and code generated in this study will be made publicly available upon the acceptance of this manuscript. The repository will include the full code architecture, PyTorch Lightning implementation, YAML configuration files, and necessary training protocols to ensure the absolute reproducibility of the reported results. Access links will be updated in the final version of the manuscript.

References↩︎