An Improved YOLOv11 for Tiny Surface Defect Detection on Electrical Commutators

Abstract

Aiming at the challenges of class imbalance, tiny defect scales, and complex brushed background interference in the surface defect detection of electrical commutators, this paper proposes a high-precision and lightweight improved instance segmentation algorithm named WG-YOLOv11. Firstly, to overcome the barrier of highly imbalanced positive and negative samples in actual industrial data collection, a Balanced Defect Synthesis (BDS) data augmentation strategy is introduced to effectively enrich the morphological diversity of tiny defects. Secondly, a Wavelet Transform Convolution (WTConv) module is collaboratively integrated into the feature extraction network to expand the receptive field while preserving the high-frequency edge details of hairline cracks. Thirdly, a Group CBAM Enhancer (GCE) module is introduced to filter out high-reflection and brushed background noise through grouped attention and weight re-calibration mechanisms. Finally, addressing the difficulty of pixel-level alignment for tiny defects, an $\alpha$-IoU loss function is utilized to improve the high-precision segmentation and localization capabilities by dynamically adjusting the gradient distribution. Comprehensive evaluations are conducted on two real-world electrical commutator surface defect datasets: KolektorSDD2 and KolektorSDD. Experimental results show that on the KolektorSDD2 dataset, compared to the YOLOv11 baseline, the Mask mAP@50 of WG-YOLOv11 increases from 85.2% to 89.2%, and the stringent metric Mask mAP@50:95 improves from 52.7% to 56.9%. Additional computational analysis on the same dataset validates that the proposed method maintains high efficiency, matching the baseline computational cost without compromising real-time inference speed. Furthermore, evaluations on the public MSD dataset confirm the model’s cross-domain generalization capabilities. The proposed framework effectively achieves a balance between detection accuracy, anti-interference robustness, and a lightweight architecture.

1. Introduction

Surface defect detection of electrical commutators (or contactors) is a crucial and challenging computer vision task in modern quality control [1,2,3]. Unlike general object detection scenarios, the surfaces of electrical commutators present strict physical and visual constraints. Firstly, defect samples are scarce, leading to an extreme class imbalance. Specifically, often only 10% of the collected data are defect samples, while the remaining 90% or more are normal samples. Secondly, the scale of defects is tiny, often appearing as low-contrast hairline cracks. Finally, abnormal areas are often hidden in complex high-reflection and brushed textures, making it difficult for models to distinguish real defects from normal structural textures [4,5,6].

Although single-stage networks like the YOLO series [7,8,9] and instance segmentation frameworks such as YOLACT [10] perform excellently in general scenarios, they exhibit distinct limitations under these industrial constraints. The limited positive samples easily cause overfitting; continuous spatial downsampling truncates high-frequency crack details; and standard regression losses lack sufficient gradient drive for pixel-level alignment of tiny targets. To address these bottlenecks, this paper proposes WG-YOLOv11, which fundamentally distinguishes itself from the standard YOLOv11 baseline by systematically replacing its general-purpose mechanisms with components tailored for industrial tiny defects. The main contributions are summarized as follows:

• Data Augmentation via Balanced Defect Synthesis (BDS): Unlike YOLOv11’s default global mixing augmentations (e.g., Mosaic) that easily destroy industrial background semantics, our BDS strategy explicitly anchors defect instances with multi-dimensional geometric variations to preserve physical realism, effectively mitigating the risk of overfitting caused by extreme data scarcity.
• Frequency-Domain Feature Extraction: Standard YOLOv11 relies on strided convolutions that irreversibly discard high-frequency topological details during spatial downsampling. We overcome this by integrating Wavelet Transform Convolution (WTConv) [11] to losslessly map features to the frequency domain, expanding the receptive field while preserving hairline crack details.
• Group Attention Re-calibration Mechanism: To address YOLOv11’s vulnerability to large-area metallic reflections, a Group CBAM Enhancer (GCE) [12] is introduced to actively filter structural background noise through a grouped re-calibration mechanism.
• Non-linear Bounding Box Regression: The default bounding box regression loss in YOLOv11 suffers from gradient flattening for tiny targets. We replace it with the $\alpha$-IoU loss [13], utilizing non-linear transformations to continuously amplify localization feedback and meet stringent high-precision segmentation requirements.

Based on the above architectural upgrades, WG-YOLOv11 is evaluated on two primary real-world industrial datasets: KolektorSDD2 and KolektorSDD. On the KolektorSDD2 dataset, its instance segmentation Mask mAP@50 reaches 89.2%, and Mask mAP@50:95 reaches 56.9%. To further verify the cross-domain robustness and real-time deployment feasibility, generalization experiments are conducted on the public MSD dataset, confirming that WG-YOLOv11 achieves a superior balance between lightweight deployment and high-precision detection compared to mainstream algorithms.

The remainder of this paper is organized as follows: Section 2 reviews related work. Section 3 details the proposed WG-YOLOv11 architecture. Section 4 presents experimental evaluations, ablation studies, and generalization analysis. Finally, Section 5 summarizes the conclusions.

2. Related Work

This section mainly reviews research in four areas closely related to this paper: surface defect detection and instance segmentation, data augmentation and sample imbalance handling, multi-scale frequency-domain features and attention mechanisms, and the evolution of bounding box regression losses.

2.1. Deep Learning-Based Surface Defect Detection and Instance Segmentation

Early industrial defect detection mainly relied on manually designed feature operators (such as Gabor filters), which had poor robustness and were susceptible to dust interference [4]. With the rise of deep learning, single-stage object detection networks represented by the YOLO series [7,8,9,14] have greatly improved the detection speed of industrial assembly lines. However, for extremely tiny and low-contrast cracks on the surface of electrical commutators, providing only bounding box detection can no longer meet high-precision industrial requirements. Therefore, instance segmentation algorithms have gradually been introduced into this field. Mainstream two-stage instance segmentation algorithms like Mask R-CNN [15] have high accuracy but slow inference speed, while single-stage algorithms like YOLACT [10] and SOLOv2 [16] improve speed but still suffer from high missed detection rates when dealing with brushed metal backgrounds. The WG-YOLOv11 proposed in this paper aims to balance high-precision pixel-level segmentation with real-time detection speed.

2.2. Data Augmentation and Sample Imbalance Handling

In actual industrial production such as electrical commutators, the imbalance between normal and defect samples is the core bottleneck restricting network performance. To solve class imbalance, Lin et al. proposed Focal Loss [17] to reduce the weight of easily classified normal samples through dynamic re-weighting. At the data level, traditional data augmentation is mostly based on global pixel operations (such as Mixup [18] or CutMix [19]), but these methods easily destroy the original background texture structure of industrial images. In recent years, the Copy-Paste instance-level augmentation algorithm proposed by Ghiasi et al. [20] provided a new idea for the long-tail distribution problem. Targeting the characteristics of micro-cracks in commutators, this paper optimizes the Copy-Paste pipeline, introduces multi-dimensional geometric variations, and proposes the BDS strategy, effectively preventing the model from overfitting to tiny positive samples.

2.3. Multi-Scale Frequency Domain Features and Attention Enhancement

To extract targets from extremely complex backgrounds, introducing attention mechanisms has become a mainstream paradigm. SENet proposed by Hu et al. [21] achieved feature re-calibration through channel compression, while CBAM proposed by Woo et al. [22] further combined spatial and channel attention to effectively suppress some background noise. However, due to the strong metal reflections and brushed textures of electrical commutators, conventional attention easily generates false alarms. On the other hand, Zhou et al. [23] demonstrated the effectiveness of Wavelet Transform in preserving high-frequency defects on industrial reflective surfaces like glass bottoms. Inspired by this, this paper innovatively combines the latest WTConv frequency-domain convolution [11] by Finder et al. with Group Cross Attention (GCE) [12], aiming to accurately extract micro-cracks and suppress background noise under complex brushed metal backgrounds.

2.4. Evolution of Bounding Box Regression Loss Functions

In object detection and instance segmentation tasks, bounding box regression loss has undergone continuous evolution. From early ${\ell }_n$-norm distances to Intersection over Union (IoU)-based losses, and later to GIoU [24], which introduces penalty terms, as well as CIoU [25], which comprehensively considers center point distance and aspect ratio, the localization accuracy of models for regular-sized targets has been significantly improved. However, for tiny defects occupying only a few pixels on the commutator surface, existing losses often lack sufficient gradient drive to achieve perfect alignment in the later stages of training. The $\alpha$-IoU proposed by He et al. [13] dynamically adjusts the gradient distribution of targets with different overlap degrees by introducing a power transformation. This paper introduces this mechanism into the tiny defect detection task and specifically optimizes the regression loss function for micro-cracks on commutators.

3. WG-YOLOv11 Defect Detection Model

To achieve accurate detection of tiny defects on electrical commutators and address the data imbalance problem in actual production, this paper proposes the WG-YOLOv11 (Wavelet and Group-attention YOLOv11) model based on the YOLOv11 baseline.

The model improvements cover three dimensions: data processing, feature extraction, and loss function optimization. Firstly, at the data level, the Balanced Defect Synthesis (BDS) strategy is introduced to alleviate sample scarcity. Secondly, the WTConv module is integrated into the backbone network to preserve the high-frequency edge details of tiny cracks. Then, the Group CBAM Enhancer (GCE) module is embedded in the neck network to filter brushed and high-reflection background noise. Finally, the $\alpha$-IoU loss function is utilized in the detection head to improve the bounding box alignment for tiny targets. The overall network architecture of the proposed WG-YOLOv11 model is illustrated in Figure 1.

3.1. Data Augmentation Strategy Based on Balanced Defect Synthesis (BDS)

Inspired by the Simple Copy-Paste algorithm [20], and addressing the scarcity of tiny defect morphologies in electrical commutators, this paper replaces traditional global pixel mixing approaches (such as Mixup or CutMix), which can disrupt industrial background textures, with an instance-level Balanced Defect Synthesis (BDS) data augmentation strategy.

Unlike the operation of directly migrating instances between random images, the BDS mechanism constructs a geometric transformation pipeline tailored for industrial defects. First, using annotation information, pixel-level defect patches ($I_1$) and their corresponding binarized instance masks (M) are extracted from the available positive sample images. Subsequently, to decouple the original defects from specific local backgrounds, geometric variations are applied to the extracted defect instances: full-angle random rotation from $0^{\circ }$∼${360}^{\circ }$ is introduced, combined with a Large Scale Jittering (LSJ) strategy. The scaling range of LSJ is set to $0.5\times$∼$1.2\times$. This decoupled transformation of scale and angle simulates imaging distortions under varying industrial lighting and depth of field, while synthesizing the phenomenon of multi-scale defects appearing in the same frame.

In the instance fusion stage, the geometrically transformed defect instances are randomly anchored and fused onto defect-free background images ($I_2$). Regarding the treatment of fusion edges, prior research indicates that the gain of complex Gaussian blending or Alpha blending on instance segmentation accuracy is negligible [20]. Therefore, to maintain data generation throughput and preserve the high-frequency edges of hairline cracks, this method omits additional edge smoothing operations and adopts a pixel-level linear combination to generate the augmented image $I_{new}$:

$$I_{new}=I_1\times M+I_2\times (1-M)$$

(1)

where $I_{new},I_1,I_2\in {\mathbb{R}}^{H\times W\times 3}$ denote the pixel matrices of the generated augmented image, the extracted defect patch, and the normal background image, respectively, with their pixel intensity values normalized to the range of $[0,1]$. $M\in {\{0,1\}}^{H\times W\times 1}$ represents the binarized instance mask, where the values are 1 for defect regions and 0 for the background; × denotes the Hadamard Product (element-wise multiplication with broadcasting).

Through this instance-level extraction and reorganization, the BDS mechanism migrates defect features to normal background canvases efficiently. This strategy encourages the network to learn the inherent texture features of the defects rather than overfitting to specific surrounding backgrounds, improving data efficiency in a few-shot environment. The complete flowchart of this instance-level BDS data augmentation is illustrated in Figure 2.

3.2. Frequency-Domain Feature Preservation Module Based on WTConv

In the surface defect detection of electrical commutators, the visual features of tiny cracks are weak and dependent on high-frequency edge gradient information. Standard convolutional architectures typically use continuous strided convolutions or pooling for spatial downsampling to expand the receptive field. However, this spatial dimension compression attenuates high-frequency details of micro-cracks, potentially leading to the loss of morphological features.

To expand the receptive field while preserving high-frequency details, a Wavelet Transform Convolution (WTConv) module [11] is integrated into the backbone network of WG-YOLOv11. WTConv utilizes the Discrete Wavelet Transform (DWT) to map features to the frequency domain for multi-band processing.

Specifically, this paper adopts the Haar wavelet basis. In the i-th layer of the cascaded structure, DWT decomposes the low-frequency input $X_{LL}^{(i-1)}$ from the previous layer into a low-frequency macroscopic structure component $X_{LL}^{\left(i\right)}$ and three high-frequency edge detail components $X_H^{\left(i\right)}=\{X_{LH}^{\left(i\right)},X_{HL}^{\left(i\right)},X_{HH}^{\left(i\right)}\}$, representing high-frequency gradients in the horizontal, vertical, and diagonal directions, respectively:

$$X_{LL}^{\left(i\right)},X_H^{\left(i\right)}=\mathrm{DWT}\left(X_{LL}^{(i-1)}\right)$$

(2)

where $X_{LL}^{(i-1)}\in {\mathbb{R}}^{C\times H_{i-1}\times W_{i-1}}$ represents the input feature map. The decomposed low-frequency component $X_{LL}^{\left(i\right)}$ and the concatenated high-frequency components $X_H^{\left(i\right)}$ both have their spatial dimensions halved and channel dimensions expanded, mapping into the domain ${\mathbb{R}}^{4C\times \frac{H_{i-1}}{2}\times \frac{W_{i-1}}{2}}$. The values of these feature elements are unconstrained real numbers ($\in \mathbb{R}$).

After the frequency domain separation, WTConv uses small-scale depthwise separable convolutions $W^{\left(i\right)}$ to perform independent feature extraction on the low-frequency and high-frequency components:

$$Y_{LL}^{\left(i\right)},Y_H^{\left(i\right)}=\mathrm{Conv}(W^{\left(i\right)},(X_{LL}^{\left(i\right)},X_H^{\left(i\right)}))$$

(3)

where $W^{\left(i\right)}$ represents the learnable weight parameters of the depthwise separable convolution at the i-th layer; $Y_{LL}^{\left(i\right)}$ and $Y_H^{\left(i\right)}\in {\mathbb{R}}^{4C\times \frac{H_{i-1}}{2}\times \frac{W_{i-1}}{2}}$ denote the output feature maps for the low-frequency and high-frequency domains, respectively.

Since DWT halves the spatial resolution of feature maps but expands the frequency channels, the same convolutional kernel can cover a larger receptive field on the original image in deeper layers. After ℓ layers of cascaded extraction, WTConv uses the Inverse Wavelet Transform (IWT) to perform a bottom-up reconstruction:

$$Z^{\left(i\right)}=\mathrm{IWT}(Y_{LL}^{\left(i\right)}+Z^{(i+1)},Y_H^{\left(i\right)})$$

(4)

where $Z^{(i+1)}\in {\mathbb{R}}^{4C\times \frac{H_{i-1}}{2}\times \frac{W_{i-1}}{2}}$ is the low-frequency aggregated feature tensor passed back from the deeper $(i+1)$-th level; and $Z^{\left(i\right)}\in {\mathbb{R}}^{C\times H_{i-1}\times W_{i-1}}$ is the reconstructed spatial-domain feature map at the current i-th layer.

Through this cascaded architecture, the receptive field grows with the number of layers ℓ, while the parameter count increases linearly, aligning with the requirements of lightweight deployment. Furthermore, by explicitly extracting high-frequency components in each layer, the edge information of micro-cracks is preserved and utilized in deeper networks, compensating for the limitations of standard convolutions.

3.3. Group CBAM Enhancer for Background Interference Suppression

Electrical commutators are processed from metal materials with complex brushed textures, which produce metal reflections under industrial lighting. In convolutional feature extraction, these structural noises often produce high activation responses, obscuring the signals of weak defects. Conventional global attention mechanisms can be dominated by high-reflection areas, causing local weak crack features to be averaged or suppressed during global pooling.

To separate defect features from complex background noise, a Group CBAM Enhancer (GCE) [12] is introduced in the neck network. GCE utilizes channel grouping and weight re-calibration to suppress background noise from local subspaces.

As shown in Figure 3, GCE first uniformly divides the input feature $\mathbf{F}\in {\mathbb{R}}^{C\times H\times W}$ into g subgroups $\{{\mathbf{F}}_1,{\mathbf{F}}_2,\dots ,{\mathbf{F}}_g\}$ along the channel dimension. A Convolutional Block Attention Module (CBAM) is then applied independently to each subgroup. The internal mechanism computes a 1D Channel Attention map (${\mathbf{M}}_c$) and a 2D Spatial Attention map (${\mathbf{M}}_s$) sequentially:

$${\mathbf{M}}_c\left({\mathbf{F}}_i\right)=\sigma \left(\mathrm{MLP}\left(\mathrm{AvgPool}\left({\mathbf{F}}_i\right)\right)+\mathrm{MLP}\left(\mathrm{MaxPool}\left({\mathbf{F}}_i\right)\right)\right)$$

(5)

$${\mathbf{F}}_i^{*}={\mathbf{M}}_c\left({\mathbf{F}}_i\right)\otimes {\mathbf{F}}_i$$

(6)

$${\mathbf{M}}_s\left({\mathbf{F}}_i^{*}\right)=\sigma \left({\mathrm{Conv}}^{7\times 7}\left([\mathrm{AvgPool}\left({\mathbf{F}}_i^{*}\right);\mathrm{MaxPool}\left({\mathbf{F}}_i^{*}\right)]\right)\right)$$

(7)

where ${\mathbf{F}}_i\in {\mathbb{R}}^{\frac{C}{g}\times H\times W}$ denotes the i-th feature subgroup. MLP, AvgPool, and MaxPool represent the Multi-Layer Perceptron, Average Pooling, and Max Pooling layers, respectively. $\sigma$ denotes the Sigmoid activation function mapping values into the range $(0,1)$, ensuring that the elements of both ${\mathbf{M}}_c\left({\mathbf{F}}_i\right)$ and ${\mathbf{M}}_s\left({\mathbf{F}}_i^{*}\right)$ strictly fall in the interval $(0,1)$. ⊗ represents element-wise multiplication with broadcasting. The local attention weights ${\mathbf{W}}_i\in {\mathbb{R}}^{\frac{C}{g}\times H\times W}$ for the subgroup are formed by:

$${\mathbf{W}}_i={\mathbf{M}}_c\left({\mathbf{F}}_i\right)\otimes {\mathbf{M}}_s\left({\mathbf{F}}_i^{*}\right),i=1,2,\dots ,g$$

(8)

where the element values of the resultant tensor ${\mathbf{W}}_i$ are strictly bounded within $(0,1)$. This grouping strategy encourages the model to mine salient features within different channel subspaces, mitigating the risk of weak defect responses being suppressed by reflection noise during global pooling.

Following the local attention computation, GCE employs an adaptive weight readjusting mechanism instead of direct feature multiplication. For each subgroup, the internal weight mean ${\mu }_i=\mathrm{mean}\left({\mathbf{W}}_i\right)$ is calculated adaptively during the network’s forward pass. To emphasize salient features and suppress background noise, the ReWeight module truncates weight elements higher than the mean to 1, while retaining the original values for elements lower than the mean, generating re-calibrated weights ${\mathbf{W}}_i^{\prime }$:

$${\mathbf{W}}_i^{\prime }(c,x,y)=\left\{\begin{array}{cc}1,\hfill & \mathrm{if}{\mathbf{W}}_i(c,x,y)>{\mu }_i\hfill \\ {\mathbf{W}}_i(c,x,y),\hfill & \mathrm{otherwise}\hfill \end{array}\right.$$

(9)

where $(c,x,y)$ are the channel, width, and height indices of the tensor; ${\mu }_i\in (0,1)$ is the scalar mean value of the local weights. Consequently, the re-calibrated weight element ${\mathbf{W}}_i^{\prime }(c,x,y)$ takes values in the range $(0,1]$.

The re-calibrated weights and the original features are combined via element-wise multiplication, and residual connections are added to generate the enhanced features ${\mathbf{F}}_i^{\prime }\in {\mathbb{R}}^{\frac{C}{g}\times H\times W}$ of the group. Finally, all subgroups are concatenated along the channel dimension to output the final enhanced feature map ${\mathbf{F}}^{\prime }\in {\mathbb{R}}^{C\times H\times W}$:

$${\mathbf{F}}_i^{\prime }={\mathbf{F}}_i\otimes {\mathbf{W}}_i^{\prime }+{\mathbf{F}}_i,i=1,2,\dots ,g$$

(10)

The ReWeight module functions as a non-linear feature equalizer. For instance, high activation weights from large reflective regions (e.g., $0.95$) and moderate weights from hairline cracks (e.g., $0.6$) are both truncated to 1 if they exceed the mean ${\mu }_i$. This truncation reduces the numerical gap, preventing structural noise from dominating the feature representation while preserving the local weights of weak defects.

3.4. Bounding Box Regression Optimization Based on $\mathsf{\alpha }$-IoU

In the detection of surface defects on electrical commutators, defect targets account for a very low proportion of pixels in the image, making the bounding box regression sensitive to coordinate offsets. When the predicted box and the ground truth box have a displacement deviation of a few pixels, the Intersection over Union (IoU) metric drops significantly.

In the standard YOLOv11 architecture, bounding box regression relies on the Complete IoU (CIoU) loss. The standard CIoU evaluates the overlap area, central point distance, and aspect ratio simultaneously, formulated as:

$${\mathcal{L}}_{\mathrm{CIoU}}=1-\mathrm{IoU}+{\displaystyle \frac{{\rho }^2(\mathbf{b},{\mathbf{b}}^{gt})}{c^2}}+\beta v$$

(11)

where $\mathrm{IoU}\in [0,1]$ is the intersection over union; $\rho (\mathbf{b},{\mathbf{b}}^{gt})$ is the Euclidean distance between the center points of the predicted box $\mathbf{b}$ and the ground truth box ${\mathbf{b}}^{gt}$; c is the diagonal length of the smallest enclosing box covering both boxes; $v\in [0,1)$ measures the consistency of the aspect ratio; and $\beta$ is a positive trade-off parameter.

While standard CIoU performs well in general, as the overlap between predicted and ground truth boxes increases during the later stages of training (i.e., $\mathrm{IoU}\to 1$), the gradient response of the linear IoU term flattens. This limits the model’s ability to achieve high-precision pixel-level alignment for tiny cracks.

To address this limitation, the $\alpha$-IoU loss function [13] is introduced to replace the default regression loss in WG-YOLOv11. This method constructs a gradient weighting mechanism by applying power transformations to the IoU term and its geometric penalties. The formulation of $\alpha$-CIoU is defined as:

$${\mathcal{L}}_{\alpha \text{-}\mathrm{CIoU}}=1-{\mathrm{IoU}}^{\alpha }+{\displaystyle \frac{{\rho }^{2\alpha }(\mathbf{b},{\mathbf{b}}^{gt})}{c^{2\alpha }}}+{\left(\beta v\right)}^{\alpha }$$

(12)

where $\alpha$ is a hyperparameter controlling the non-linear power transformation. The non-linear operation is introduced by raising the IoU term, the normalized central distance term ${\displaystyle \frac{{\rho }^2}{c^2}}$, and the aspect ratio penalty term $\beta v$ to the power of $\alpha$ simultaneously.

The partial derivative of the $\alpha$-IoU loss with respect to the IoU term is calculated as:

$${\displaystyle \frac{\mathsf{\partial }{\mathcal{L}}_{\alpha \text{-}\mathrm{CIoU}}}{\mathsf{\partial }\mathrm{IoU}}}=-\alpha ·{\mathrm{IoU}}^{\alpha -1}$$

(13)

This partial derivative reflects the gradient magnitude backpropagated by the model during training. To achieve gradient up-weighting for high-IoU samples, $\alpha$ is set strictly greater than 1. Since $\mathrm{IoU}\in [0,1]$ and $\alpha >1$, the gradient magnitude $\left|{\displaystyle \frac{\mathsf{\partial }\mathcal{L}}{\mathsf{\partial }\mathrm{IoU}}}\right|=\alpha ·{\mathrm{IoU}}^{\alpha -1}$ monotonically increases with the IoU value.

The exponential power term functions as an adaptive regulator for the gradient magnitude. For predicted boxes with high IoU, the non-linear transformation assigns higher gradient weights, generating stronger feedback signals during backpropagation. For commutator cracks, this non-linear weighting mitigates the gradient flattening issue, encouraging the network to focus on fine-tuning minor offsets and improving boundary alignment.

4. Experimental Results and Analysis

4.1. Experimental Datasets and Evaluation Metrics

This paper conducts comprehensive validations on two primary real-world industrial surface defect datasets: KolektorSDD2 and KolektorSDD.

The KolektorSDD2 dataset is one of the most challenging benchmark datasets in the field of industrial defect detection. Its original dataset is divided into a training set and a validation set. The original training set contains 2331 images (including 246 defective positive samples and 2085 defect-free negative samples), and the original validation set contains 1004 images (including 110 positive samples and 894 negative samples). To address the severe class imbalance in the original data, this paper utilizes the Balanced Defect Synthesis (BDS) strategy to perform instance-level expansion synchronously on both sets.

To maximize the morphological diversity of the synthesized defects and avoid model overfitting to specific discrete geometric transformations, the BDS strategy employs a continuous random sampling mechanism for its design parameters. Specifically, during the synthesis process, the rotation angle is randomly drawn from a full-angle uniform distribution $\theta \sim U(0^{\circ },{360}^{\circ })$, and the scale jittering factor is randomly sampled from a uniform distribution $S\sim U(0.5,1.2)$. This continuous random variation covers an infinite combinatorial space of defect morphologies within reasonable physical boundaries. Through this strategy, the total number of images in the augmented training set is expanded from 2331 to 9294, where the positive samples synthesized via the BDS strategy are increased to 2562, and the positive sample ratio is corrected from the original 10.55% to 27.57%, significantly alleviating the class imbalance problem. It is worth noting that while validation sets are typically kept unchanged, the original validation set contained extremely scarce positive samples (only 110 defects, retaining the skewed 1:9 imbalance ratio). Evaluating stringent instance segmentation metrics on such a limited pool leads to severe statistical variance. Therefore, the BDS augmentation was cautiously applied to the validation set exclusively to increase the volume of positive evaluation instances (to 223 defect samples). This ensures a statistically robust and reliable assessment of the model’s generalization capability against diverse defect morphologies. Strict data isolation was maintained between the training and validation synthesis pools to prevent any data leakage. The specific data distribution comparison for KolektorSDD2 is shown in

Table 1. Detailed composition comparison of KolektorSDD2 surface defect dataset before and after augmentation.

Data Stage	Subset Type	Total Images	Defective	Defect-Free	Defect Ratio
Original Stage	Train Set	2331	246	2085	10.55%
	Val Set	1004	110	894	10.96%
BDS Augmented	Train Set	9294	2562	6732	27.57%
	Val Set	2200	223	1977	10.14%

.

In addition to KolektorSDD2, this paper also introduces the KolektorSDD dataset for comparative evaluation. This dataset contains a total of 399 high-resolution images, of which 52 are defect samples. Since the KolektorSDD dataset is mainly used to verify the model’s generalization performance under strong reflection and complex brushed backgrounds, no additional augmentation operations were performed on it in the experiments, and its original distribution was kept for evaluation.

In terms of evaluation metrics, to comprehensively evaluate the detection and segmentation performance of the model, this paper adopts the conventional Precision (P) and Recall (R) as fundamental metrics, and introduces the mean Average Precision (mAP) for comprehensive evaluation. The Average Precision (AP) of each category can be obtained by calculating the area under the Precision–Recall (P-R) curve, and then taking the average over all N categories to obtain the mAP. Its basic formula is as follows:

$$mAP={\displaystyle \frac{1}{N}}{\displaystyle \sum _{i=1}^N}{\int }_0^{1}P\left(R\right)dR$$

(14)

The conventional mAP@50 refers to the mAP value calculated when the Intersection over Union (IoU) threshold between the predicted box and the ground truth box is set to 0.5, which is mainly used to measure the model’s basic object localization and recognition ability. Furthermore, to verify the high-precision pixel-level alignment capability of the loss function for tiny defect bounding boxes, this paper specially introduces a more stringent comprehensive evaluation metric, mAP@50:95. This metric refers to performing 10 mAP calculations with a step size of 0.05 in the IoU threshold range from 0.50 to 0.95 and taking the average. Its calculation formula is briefly described as follows:

$$mAP@50:95={\displaystyle \frac{1}{10}}\sum _{t\in \{0.50,0.55,\dots ,0.95\}}{mAP|}_{IoU=t}$$

(15)

Compared with a single threshold, mAP@50:95 can acutely reflect the localization penalties under high-stringency overlap requirements, and is the core metric for verifying the detection accuracy of tiny targets.

4.2. Experimental Environment and Training Setup

All experiments were conducted on a single NVIDIA GeForce RTX 4090 GPU (24 GB VRAM) using Python 3.10.19, PyTorch 2.10.0, and the Ultralytics 8.4.19 framework. For the network architecture, the baseline model is the standard YOLOv11n-seg loaded with official pre-trained weights. Our proposed WG-YOLOv11 modifies the YOLOv11-seg codebase by integrating the WTConv, GCE, and $\alpha$-IoU modules.

The network input size was set to $640\times 640$ with a batch size of 16. Models were trained for a maximum of 1000 epochs, incorporating an early stopping patience of 100 epochs to prevent overfitting. Parameter optimization was performed using the Stochastic Gradient Descent (SGD) optimizer. The learning rate was set to 0.01, with a momentum of 0.937 and a weight decay of 0.0005. To stabilize the initial gradient updates, a 3-epoch warm-up strategy was applied. For data augmentation, the Mosaic probability was set to 1.0 and disabled during the final 10 epochs. To ensure strict reproducibility, the global random seed was fixed at 0, CuDNN deterministic execution was enforced, and Automatic Mixed Precision (AMP) was enabled. The core hyperparameter settings are summarized in

Table 2. Core training hyperparameters and optimization settings.

Hyperparameter	Value	Hyperparameter	Value
Input Size	$640\times 640$	Batch Size	16
Max Epochs	1000	Early Stopping	100
Optimizer	SGD	Learning Rate ($l{r}_0$)	0.01
Momentum	0.937	Weight Decay	0.0005
$\alpha$ factor	3.0	Warm-up Epochs	3

.

4.3. Ablation Study Design and Analysis

To verify the individual impact and the combined effects of the proposed Balanced Defect Synthesis (BDS) strategy, WTConv module, GCE mechanism, and $\alpha$-IoU loss function, a comprehensive ablation study was conducted on the KolektorSDD2 dataset. Addressing the need to isolate the contributions of each component, we evaluated the modules both independently and incrementally. Since the core task involves the extraction of tiny defects, the evaluation focuses on instance segmentation (Mask) metrics. The detailed ablation design combinations are shown in

Table 3. Comprehensive ablation study design for electrical commutators on KolektorSDD2.

Method	Baseline	BDS	WTConv	GCE	$\mathsf{\alpha }$-IoU ($\mathsf{\alpha }=3$)
1	√
2	√	√
3	√		√
4	√			√
5	√				√
6	√		√	√	√
7	√	√	√
8	√	√	√	√
9	√	√	√	√	√

, and the corresponding quantitative results are presented in

Table 4. Comparison results of instance segmentation performance in ablation experiments.

Method	Mask P/%	Mask R/%	Mask mAP@50/%	Mask mAP@50:95/%
1 (Baseline)	81.70	45.10	59.10	24.10
2	84.80	79.80	85.20	52.70
3	80.90	63.90	73.90	33.70
4	77.40	63.20	68.00	30.50
5	76.30	70.30	74.90	35.10
6	75.30	66.60	73.10	32.10
7	90.00	78.90	87.10	54.40
8	89.00	83.20	88.00	56.90
9 (Ours)	88.00	85.70	89.20	56.90

.

The ablation results indicate that the baseline model (Method 1) yields a Mask Recall of 45.10% and a mAP@50 of 59.10%. Applying the BDS strategy alone (Method 2) provides a substantial performance increase, raising the Recall to 79.80% and the mAP@50 to 85.20%. This confirms the necessity of enriching morphological diversity to overcome the feature learning barrier caused by extreme data scarcity.

To explicitly isolate the contributions of the network architecture, we evaluated the structural and loss modules independently without the BDS strategy (Methods 3 to 6). Incorporating WTConv (Method 3) increases the mAP@50 to 73.90%, reflecting its ability to extract high-frequency details. Adding GCE (Method 4) improves the mAP@50 to 68.00%, demonstrating its effect in filtering background noise. Applying the $\alpha$-IoU loss function (Method 5) raises the Recall to 70.30% and mAP@50 to 74.90%, suggesting that the adjusted gradient weighting assists in bounding box alignment. However, combining all three modifications without data augmentation (Method 6) results in a mAP@50 of 73.10%. Compared to Method 2, this clearly indicates that without addressing the initial data scarcity issue, structural modifications alone have limited effectiveness and inevitably face severe overfitting.

Building upon the robust data foundation provided by BDS, the sequential integration of the architectural modules further refines the detection results. The addition of WTConv (Method 7) raises the Mask Precision to 90.00%, exhibiting strong anti-false alarm capabilities. Including GCE (Method 8) increases the stringent mAP@50:95 metric to 56.90%. Finally, incorporating the $\alpha$-IoU loss function to form the complete WG-YOLOv11 (Method 9) yields the highest Mask mAP@50 of 89.20% and a peak Recall of 85.70%. This step-by-step evaluation demonstrates a highly functional integration of the proposed data processing, feature extraction, and regression optimization approaches.

Sensitivity Analysis of Hyperparameter $\alpha$

To determine an appropriate non-linear scaling factor for the $\alpha$-IoU loss function, a parameter sensitivity analysis was conducted for $\alpha \in \{2,3,4,5\}$ based on the Method 8 setup. The comparative results are summarized in

Table 5. Sensitivity analysis of the hyperparameter $\alpha$ in $\alpha$-IoU (KolektorSDD2 dataset).

$\mathsf{\alpha }$ Value	Mask P/%	Mask R/%	Mask mAP@50/%	Mask mAP@50:95/%
$\alpha =2$	88.10	81.20	87.10	56.40
$\alpha =3$	88.00	85.70	89.20	56.90
$\alpha =4$	88.40	82.20	87.90	56.10
$\alpha =5$	87.70	79.90	86.10	55.10

.

As shown in Table 5, the segmentation performance varies with the choice of $\alpha$. When $\alpha$ is set to 2, the gradient weighting modification provides limited improvement for boundary alignment compared to higher values. Conversely, larger values such as $\alpha =4$ or 5 introduce larger gradient penalties to high-IoU samples, which can lead to optimization instability and decrease both the Recall and overall mAP metrics. The model achieves its best overall results at $\alpha =3$. This configuration adjusts the gradient distribution to facilitate the fitting of micro-crack boundaries, leading to the final performance of the WG-YOLOv11 model.

4.4. Performance Comparison with Mainstream Segmentation Models

To comprehensively verify the performance advantages of the proposed WG-YOLOv11 in different industrial environments, this paper executes mainstream algorithm comparison experiments on both the KolektorSDD2 and KolektorSDD datasets. Participating models include classic semantic segmentation networks UNet [26] and DeepLabv3+ [27], the mainstream single-stage instance segmentation network YOLOv8n-seg [8], the original baseline model YOLOv11, as well as the recently introduced YOLOv12 [28] and YOLOv26 [29] architectures. Comprehensive performance evaluation results for different models on both datasets are shown in

Table 6. Performance comparison of different segmentation models on KolektorSDD2 and KolektorSDD datasets.

Dataset	Model Name	Mask P/%	Mask R/%	Mask mAP@50/%
KolektorSDD2	UNet [26]	36.00	32.64	29.18
	DeepLabv3+ [27]	35.88	33.90	30.41
	YOLOv8n-seg [8]	82.00	72.80	77.50
	YOLOv11	84.80	79.80	85.20
	YOLOv12 [28]	91.30	83.40	88.90
	YOLOv26 [29]	90.90	81.00	86.70
	WG-YOLOv11	88.00	85.70	89.20
KolektorSDD	UNet [26]	44.14	36.69	32.89
	DeepLabv3+ [27]	39.06	33.33	29.40
	YOLOv8n-seg [8]	89.00	90.00	94.20
	YOLOv11	80.30	100.0	94.80
	YOLOv12 [28]	78.40	80.00	80.80
	YOLOv26 [29]	100.0	79.90	92.00
	WG-YOLOv11	96.40	100.0	99.50

.

From the KolektorSDD2 comparison data in Table 6, it can be observed that traditional fully convolutional networks (UNet and DeepLabv3+) experience feature degradation during continuous downsampling when processing fine cracks, resulting in mAP@50 scores of approximately 30%. In contrast, the YOLO series demonstrates more stable baseline performance. When compared to the recently introduced YOLOv12 and YOLOv26 models, the proposed WG-YOLOv11 shows specific performance characteristics. Although YOLOv12 and YOLOv26 achieve slightly higher Mask Precision (91.30% and 90.90%, respectively), WG-YOLOv11 attains the highest Mask Recall (85.70%) and the highest overall Mask mAP@50 (89.20%). In industrial quality control, a high Recall rate is critical to minimize false negatives and prevent defective products from passing inspection. Facilitated by the retention of high-frequency edge features via the WTConv module and the boundary alignment provided by the $\alpha$-IoU loss, WG-YOLOv11 improves upon the original YOLOv11 baseline by 4.0% in mAP@50. This indicates its effectiveness in extracting and recognizing tiny defect features compared to the evaluated alternative architectures.

Furthermore, on the KolektorSDD dataset, which features highly reflective surfaces, traditional semantic segmentation models yield accuracy metrics below 35%, partly due to limitations in suppressing reflection noise. The YOLOv12 and YOLOv26 models exhibit an imbalance between Precision and Recall in this scenario. While YOLOv26 achieves a Mask Precision of 100.0%, its Mask Recall decreases to 79.90%, indicating that approximately 20% of the actual defects are missed—a substantial limitation for industrial defect detection. The proposed WG-YOLOv11 offers a more balanced performance in this context. Compared to the YOLOv11 baseline, WG-YOLOv11 increases Mask Precision from 80.30% to 96.40% while maintaining a 100.0% Recall rate, resulting in a Mask mAP@50 of 99.50%. Additionally, the Mask mAP@50:95 metric, which reflects pixel-level alignment accuracy, achieves a stable increase of 2.5%. This demonstrates that the integration of the GCE and WTConv modules helps WG-YOLOv11 reduce false positives and maintain robustness against complex industrial backgrounds, offering an improvement over the baseline and other evaluated YOLO variants. Visual comparison results are shown in Figure 4 and Figure 5.

4.5. Computational Efficiency and Generalization Analysis

In industrial quality control, the defect detection model must balance high detection accuracy with the requirements for real-time processing and lightweight deployment. This section provides a detailed analysis of the computational cost, inference speed, and model complexity of the proposed WG-YOLOv11 compared to the baseline YOLOv11 model.

To ensure the statistical stability of the speed evaluation, metrics for model complexity and pure network inference speed were benchmarked over the KolektorSDD2 validation set. To explicitly demonstrate the computational impact of each proposed component, we evaluated the independent performance of WTConv, GCE, and the $\alpha$-IoU loss function. The detailed comparison results are summarized in

Table 7. Comparison of model complexity, computational cost, and inference speed.

Model	Params (M)	GFLOPs	Inference (ms)	*Total Time (ms)
YOLOv11	2.84	9.7	0.9	4.9
WTConv Only	2.78	9.6	1.0	2.3
GCE Only	2.98	9.8	1.1	2.6
$\alpha$-IoU Only	2.84	9.7	0.8	1.5
WG-YOLOv11	2.92	9.7	1.1	4.8

*Total Time includes pre-processing, pure network inference, and NMS post-processing.

.

Model complexity is primarily measured by the number of learnable parameters (Params), while the computational cost is evaluated using Giga Floating-point Operations (GFLOPs). As shown in Table 7, the baseline model possesses 2.84 M parameters and 9.7 GFLOPs. Substituting standard spatial convolutions with WTConv independently reduces the parameters and computational cost to 2.78 M and 9.6 GFLOPs, owing to the lightweight nature of depthwise separable convolutions. Conversely, the GCE module introduces a slight increase in complexity (2.98 M parameters and 9.8 GFLOPs) due to its multi-branch attention structure. Notably, as the $\alpha$-IoU is a loss function modification used during training, it maintains the same parameters and GFLOPs as the baseline during the inference phase.

When these modules are integrated into the final WG-YOLOv11 architecture, they achieve a favorable structural balance. The parameter count is slightly adjusted to 2.92 M, while the overall computational cost remains consistent with the baseline at 9.7 GFLOPs. This indicates that the proposed framework enhances feature extraction and noise suppression without imposing significant additional computational overhead, making it suitable for edge devices in industrial environments.

To further evaluate the cross-domain generalization capability and computational efficiency of the proposed method on multi-class industrial targets, the Mobile phone screen surface defect (MSD) dataset is introduced. The MSD dataset, built by Peking University, contains 1200 high-resolution ($1920\times 1080$) images captured by industrial cameras, encompassing three distinct types of surface defects: Oil, Scratch, and Stain (400 images for each type). The dataset is randomly divided into training, validation, and testing sets in a 6:2:2 ratio. To verify the intrinsic generalization of the network architecture across different industrial scenes, the original data distribution of the MSD dataset was maintained for evaluation.

The comparative results on the MSD dataset are presented in

Table 8. Comparison of instance segmentation performance on the MSD dataset.

Dataset	Model Name	Mask P/%	Mask R/%	Mask mAP@50/%
MSD	YOLOv11	97.40	97.70	97.50
	WG-YOLOv11	98.10	97.90	97.80

. WG-YOLOv11 improves overall detection stability across different defect types, raising Mask Precision from 97.40% to 98.10%, Mask Recall from 97.70% to 97.90%, and Mask mAP@50 from 97.50% to 97.80%. In addition to the accuracy improvements, the proposed structural optimizations also demonstrated excellent computational efficiency. The pure network inference speed on the high-resolution MSD dataset was reduced from 1.7 ms per image (baseline) to 0.9 ms per image (WG-YOLOv11). These comprehensive evaluations indicate that the architectural innovations of WG-YOLOv11 not only extract multi-class defect features stably but also achieve a superior balance between detection accuracy, cross-domain generalization capability, and inference speed.

5. Conclusions

To address the limitations in detecting tiny surface defects during industrial inspection—specifically those arising from extreme sample imbalance, small defect scales, and complex background interference—this paper presents WG-YOLOv11, an improved instance segmentation network.

The framework introduces several targeted optimizations. Methodologically, the Balanced Defect Synthesis (BDS) data augmentation strategy is employed to mitigate feature learning barriers associated with data scarcity. Structurally, integrating the WTConv module into the backbone facilitates the preservation of high-frequency edge information of micro-cracks through frequency-domain feature processing. To reduce the impact of metallic reflections, the GCE module is embedded in the neck network for attention re-calibration, assisting in the suppression of structural background noise. Finally, the $\alpha$-IoU loss function is utilized in the regression branch to enhance pixel-level boundary alignment by adjusting the gradient weighting for high-overlap samples.

Experimental results on two real-world electrical commutator datasets indicate that WG-YOLOv11 provides consistent improvements in instance segmentation performance. On the KolektorSDD2 dataset, the proposed WG-YOLOv11 increases the Mask mAP@50 from 85.2% to 89.2% and achieves a Mask mAP@50:95 of 56.9%. On the highly reflective KolektorSDD dataset, WG-YOLOv11 increases the Mask Precision to 96.40% while strictly maintaining a perfect 100.0% Recall rate, achieving a near-perfect Mask mAP@50 of 99.50%. Furthermore, cross-domain evaluations on the MSD dataset confirm the model’s generalization across different industrial scenarios, maintaining stable segmentation accuracy and accelerating inference speeds. By balancing detection robustness with a lightweight architecture, the proposed method offers a practical reference for automated industrial quality control systems.

Despite these promising results, the current study acknowledges certain limitations. First, regarding the BDS data augmentation strategy, employing continuous random sampling prevents the quantitative ablation of discrete geometric parameters, which limits the identification of an optimal static configuration. Second, while WG-YOLOv11 effectively manages structural brushed textures and typical metallic reflections, its generalization capability may be challenged under extreme, unpredictable random lighting interference or highly complex non-structural noise (e.g., severe oil stains randomly overlapping with specular halation). In such extreme optical environments, the weak gradient features of hairline cracks might be completely overwhelmed, potentially leading to missed detections or false positives. Therefore, future work will focus on exploring adaptive augmentation parameter search techniques to optimize synthesis bounds quantitatively, and integrating optical-invariant feature learning to further enhance the model’s robustness across highly unpredictable industrial scenarios.