WTAYOLO: A Detection Model for Weak Scratches on Metal Surfaces

The problem

Scratches on lithium-battery surfaces affect the batteries’ electrochemical performance and safety in operation, so they must be found during quality control. On highly reflective metal this is hard. The aluminium alloy reflects light strongly and is very sensitive to lighting, while shallow scratches have little depth, blurred contours and minimal contrast with the surrounding surface. Scratches are also randomly placed and irregularly shaped, and many look much like the surface itself.

Photometric stereo, which images the surface under light from several directions, gives detailed surface topography without complex laser equipment, so it suits fine scratches on reflective metal. On the algorithm side, U-Net-style networks are widely used for defect segmentation, and YOLO networks are fast and accurate. According to the authors, few studies have applied YOLO segmentation to metal scratches, and balancing accuracy, model size and real-time performance remains difficult. This paper builds a scratch segmentation model, WTAYOLO, on the YOLOv11n baseline and feeds it photometric stereo images.

How it works

WTAYOLO keeps the structure of YOLOv11n and replaces its C3K2 modules with C3K2-WTAConv. Inside each bottleneck, a new convolution, WTAConv, combines wavelet-transform convolution (WTConv) with the CBAM attention module.

1. Photometric stereo images. The metal surface is imaged by a photometric stereo system, which reduces interference from surface reflections and brings out the scratches.
2. Wavelet convolution for a large receptive field. WTConv splits its input with a 2D Haar wavelet transform into a low-frequency and three high-frequency components, applies small convolution kernels to each, and reconstructs the output with the inverse transform. Repeated over two levels, this gives a large receptive field while the parameters grow only logarithmically, avoiding the cost of large kernels.
3. Channel and spatial attention (CBAM). The output of the wavelet convolution passes through CBAM. Channel attention pools the feature map globally (average and max), passes both descriptors through a shared MLP and weights each channel; spatial attention pools along the channel axis, applies a 7×7 convolution and weights each location. This strengthens important features and suppresses background.
4. Segmentation. The modified backbone and neck feed YOLOv11’s segmentation head, which outputs scratch masks.

Results

The test rig is a photometric stereo system built from an OPT camera (IMX183 sensor, 5472×3648), a 35 mm OPT lens and an OPT ring light (Table 1). With it the authors captured 566 images of scratched metal sheets at 1200×1200 pixels, annotated at pixel level by experts, and split them into 510 training and 56 validation images. All models were trained with the same parameters.

The comparison covers Otsu thresholding, U-Net, YOLOv5n, YOLOv8n and YOLOv11n. In the authors’ examples, Otsu fails to suppress background noise, so the scratches are obscured; U-Net partly overcomes the background but misses low-contrast scratches, leaving incomplete results; the YOLO networks suppress the background and use multi-scale information, and WTAYOLO segments the weak scratches best. WTAYOLO also has the highest precision (0.942) and recall (0.736).

On size, WTAYOLO has 16.3% fewer parameters than YOLOv8n (3.258 M) and 3.8% fewer than YOLOv11n. Adding the attention costs only 0.003 M parameters. By comparison, the larger YOLOv11s reaches 0.841 mAP50 with 10.067 M parameters and 35.3 GFLOPs. The extra modules do cost time: WTAYOLO needs 20.9 ms per image against 14.5 ms for YOLOv11n (Table 3). The authors judge that this still meets the time constraints of industrial inspection, and name optimising the algorithm for real-time detection as future work.