An Enhanced YOLOv5-Based Algorithm for Metal Surface Defect Detection

Abstract

The detection of surface defects in metal materials has been a challenging issue in the industrial domain. The existing algorithms for metal surface defect detection are limited to a few specific types of defects and exhibit low performance with detection of defects of varying scales. A novel detection method based on the Information Enhancement YOLOv5 Network (IE-YOLOv5) for surface defects in metal parts is proposed, to realize efficient detection, which introduces a lightweight Federated Fusion Slim Neck module (FF-Slim-Neck) and a Parameter-free Spatial Attention mechanism (PSA) in YOLOv5 network. Comparative experiments were conducted using the NEU-DET dataset. The experimental results indicate that the proposed algorithm for detecting defects on metal surfaces achieves an average precision of 96.7% when identifying six different types of surface imperfections: crazing, inclusions, patches, pitting, scaling, and scratches. This represents a 2.4% enhancement in precision compared to the YOLOv5 algorithm. The measured processing velocity of this approach stands at 46.17 frames per second (FPS), highlighting its remarkable qualities of resilience, precision, and real-time capability.

1. Introduction

The identification of surface irregularities on metal workpieces plays a pivotal role in guaranteeing the quality of industrial manufacturing during the production process. During the fabrication and manufacturing of metal workpieces, factors such as machining methods and manufacturing conditions can result in the emergence of imperfections like scratches, protrusions, and abrasions on the surface of material. These defects exert a substantial influence on the quality and functionality of the products, diminishing their longevity and potentially leading to the production mishaps in more serious instances. Manual screening is typically employed for detecting metal surface defects in industrial production currently. Nevertheless, this method is marked by elevated expenses, reduced efficiency, and the quality of detection is significantly affected by factors like the skill level of workers, posing challenges in fulfilling the requirements of large-scale product inspection.

There are two main categories of methods for metal surface defect detection [1]. One relies on conventional image processing methodologies, and the other is rooted in deep learning techniques.

Ren et al. [2] introduced a similarity assessment approach reliant on color histograms to evaluate the Electromagnetic Resonance Tomography (ERT) imagery of resistor blocks. An Adaptive Cuckoo Search (ACS) algorithm is used to compute the optimal threshold for multi-level grayscale image thresholding improvement. Chen et al. [3] used a Genetic Algorithm (GA) to calculate the optimal Gabor filters to suppress texture interference. A novel fabric defect detection method is proposed, which used ACS-based thresholding and GA-based optimal Gabor filters. Wang et al. [4] introduced a technique for detecting color-based point defects. This approach relies on factors such as local contrast analysis, point defect enhancement, noise reduction, and background suppression, all in conjunction with the utilization of the OR operation.

Defect detection methods based on deep learning primarily encompass the Faster R-CNN, various networks from the YOLO series, and the SSD network. Bayraktar et al. [5] leveraged the strengths of GANs while addressing the limitations of conventional approaches, ultimately improving the overall performance of surface defect detection. Xia et al. [6] proposed an innovative defect detection approach utilizing Faster R-CNN specifically designed for identifying plate surface imperfections. A bilateral filtering algorithm was used, in which a shape-adaptive convolution with ResNet50 was applied to obtain defect semantic feature mapping in the feature pyramid network. Zhang et al. [7] introduced a surface defect detection approach utilizing MobileNet-SSD as the foundation. They streamlined the detection model by making adjustments to both the network structure and parameters. Zheng et al. [8] improved the YOLOv4-Tiny algorithm by employing a parallel feature processing strategy to extract features from bimodal images. Features from different layers were fused to enhance the defect feature extraction approach.

However, the ACS algorithm may require a large number of iterations to find a satisfactory optimization solution, especially for complex optimization problems. This may result in long runtimes. Faster R-CNN only employs features extracted from the final convolutional layer of the underlying network for detection. This deficiency in incorporating shallow features results in subpar performance when identifying minor defects in small targets, and it also results in a relatively bulky model that hampers detection speed. Although the SSD algorithm predicts smaller targets using shallow features, it still lacks sufficient deep feature information, resulting in limited capability in detecting small targets. The YOLOv4-Tiny algorithm has faster detection speed but requires improvement in accuracy.

Facing the challenges of general object detectors and their backbone networks in small object defect detection tasks, as well as the difficulty in balancing the detection accuracy and efficiency of existing metal surface defect detection, a surface defect detection network called the Information Enhancement YOLOv5 network (IE-YOLOv5) is proposed in this paper. The standard convolution block in the neck module is replaced with a lightweight one in the IE-YOLOv5 network to construct a lightweight Neck module-GSConv, which can enhance the feature extraction efficiency. A slim neck structure with feature fusion based on the federated learning algorithm (FedFusion) [9], FF-Slim-neck, is proposed to optimize the Neck module. The Parameter-free Spatial Attention mechanism (PSA) is introduced into the YOLOv5 network to enhance the feature extraction capability, enabling swift and precise detection of defects on metal surfaces.

In summary, the contributions of this study are as follows:

(1)

The novel lightweight convolution module GSConv is employed to replace the standard convolution in the original YOLOv5 network’s Neck module, which allows efficient preservation of hidden links for each channel and feature extraction efficiency improvement.

(2)

Utilizing the federated learning algorithm known as FedFusion, we introduce a feature fusion module designed to aggregate features from both local and global models at the feature extraction stage. Combined with GSConv, the FF-Slim-Neck structure is constructed to optimize the Neck module of the network and reduce additional computational costs.

(3)

A lightweight Parameter-free Spatial Attention mechanism called PSA is proposed to reduce network parameters and enhance network representation capacity. Additionally, the YOLOv5’s CIoU loss function is optimized to EIoU loss function, which improves prediction accuracy and speeds up model convergence.

2. Methods

2.1. Network Architecture of YOLOv5

The YOLOv5 network consists primarily of four components: Input, Backbone, Neck, and Head. Figure 1 presents an illustration of the holistic structure of YOLOv5. The input is a preprocessed 640 × 640 × 3 image. CBL represents the basic convolutional unit; CSP_x denotes the feature extraction module; Concat represents the concatenation module; upsample denotes the upsampling module; and SPP refers to the Spatial Pyramid Pooling layer for efficient spatial pooling; The dimensions of the output feature maps are 20 × 20 × 255, 40 × 40 × 255, and 80 × 80 × 255, representing the length, width, and depth, respectively. These dimensions correspond to the network’s detection tasks for large, medium, and small objects of different scales.

2.2. Network Architecture of IE-YOLOv5

In the original YOLOv5 [10], the input image is processed through convolution and upsampling in the Neck module, which results in the gradual transfer of spatial information to the channels. When the feature map’s spatial dimensions are compressed and the channels are expanded, there is a partial loss of semantic information. The GSConv structure can be introduced into the Neck module to address this issue.

We introduce the FF-Slim-Neck to strike a balance among detection accuracy, detection speed, and computational cost within GSConv. This addition aids in addressing the challenge of small object detection while upholding overall efficiency.

To further enhance the performance of the detection algorithm, we have implemented various improvements in image feature extraction, feature aggregation, fusion, and network lightweight design. The original FPN + PAN feature fusion network in the Neck structure is replaced with FF-Slim-Neck and the GSConv combined feature fusion network. This novel design enhances the network’s capacity for crucial information extraction while also optimizing the efficiency of feature fusion.

Additionally, the Head layer incorporates the PSA attention mechanism using the concept of three-dimensional weights to assess statements. This results in a reduction in inference time and an enhancement of the network’s detection speed.

The improved algorithm framework IE-YOLOv5 based on YOLOv5 is shown in Figure 2. These modifications collectively enhance the detection algorithm’s performance, particularly in detecting small objects in the industrial environment.

2.3. GSConv Structure

In the original YOLOv5 [11], the input image is processed through convolution and upsampling in the Neck module, which results in the gradual transfer of spatial information to the channels. When the feature map’s spatial dimensions are compressed and the channels are expanded, there is a partial loss of semantic information. The GSConv structure can be introduced into the Neck module to address this issue.

GSConv is a lightweight convolutional layer design that significantly reduces the number of parameters and the computational complexity of the model. With GSConv, IE-YOLOv5 is able to reduce the computational requirements of the model while maintaining good performance, thus improving the efficiency of feature extraction. The GSConv can preserve hidden connections for each channel efficiently. Given an input feature map of size C₁ × H × W, the GSConv module applies two types of convolutions: standard Convolution (SC) and depth-wise separable convolution (DSC) [12]. SC is applied to conduct convolution on individual groups within the input feature map, utilizing a convolution kernel sized C₂/2 × C₁ × h₁ × w₁, yielding a feature map of dimensions C₂/2 × H × W. DSC is then used for feature map segmentation and channel fusion, producing an output feature map with channels still at C₂/2 × H × W. The feature maps from SC and DSC are fused using the Concat operation to obtain a feature map of dimensions C₂ × H × W. Finally, the Shuffle module is employed for feature reorganization, yielding an output feature map of dimensions C₂ × H × W, as depicted in Figure 3. The GSConv module facilitates the propagation of information generated by SC into all segments of the information produced by DSC. The processed image features contain less redundant and repetitive information, reducing the parameter count and floating-point operations, thereby improving the efficiency of feature extraction.

2.4. FF-Slim-Neck Structure

The depth of the network is increased with the introduction of GSConv, which increases the detection time significantly as shown in Figure 4. p1, p2, and p3 refer to the images whose sizes are 128 × 80 × 80, 256 × 40 × 40, 512 × 20 × 20 after the original images are processed by Backbone. Deeper networks can introduce more resistance to data flow, significantly increasing detection time. To address this issue, the design of the FedFusion Slim Neck (FF-Slim-Neck) structure is inspired by generalized approaches to CNN learning capabilities found in methods like DensNet [13], VoVNet [14], and CSPNet [15].

The FF-Slim-Neck structure incorporates GSConv, which may balance the trade-off between network depth and efficiency. This design allows the network to extract important features while keeping the computational cost and detection time at a reasonable level. The FF-Slim-Neck structure combines the advantages of GSConv with the principles of these generalized CNN learning methods to achieve improved performance in detecting metal part defects.

Local features pertain to data extracted from relatively compact or specific regions within the dataset. These features are generally employed to capture details and fluctuations in the data, focusing on local structures, and do not take into account the broader context of the entire dataset. In the case of images, local features might encompass small segments of an image, like edges, corners, or texture patterns. In contrast, global features encompass the information gathered from the entirety of the dataset or the entire data range. They are utilized to grasp the overarching context and global relationships within the data, facilitating a more comprehensive comprehension of the entire dataset.

Incorporated into the FCConv model, as depicted in Figure 5, is the Federated Learning algorithm known as FedFusion [16], which is based on feature fusion techniques. The input image X undergoes transformation into two feature spaces, with the local feature extractor represented as E_l and the global feature extractor as E. This transformation results in the generation of feature mappings, denoted as E_l(x), E_g(x) ∈ R^C×H×W. Subsequently, E_l and E_g are embedded into the fusion feature space using a 1 × 1 convolution. The utilization of 1 × 1 convolutions serves as an effective strategy for dimensionality reduction within feature channels, imparting significant advantages in terms of computational complexity and model parameter count reduction. This approach effectively mitigates computational and storage requirements.

$$F_{conv}\left(E_l\left(x\right),E_g\left(x\right)\right)=W_{conv}\left(E_l\left(x\right)| |E_g\left(x\right)\right)$$

(1)

$W_{conv}\in R^{2C\times C}$ represents the learned weight matrix, and || denotes the operation of concatenating feature maps along the channel axis.

The GSConv and FedFusion are utilized in the Neck module to design the FF-Slim-Neck structure. The design of the cross-stage partial network FC module employs a one-step aggregation technique. This approach effectively lowers the computational complexity and reduces the inference time during the detection process. The FF-Slim-Neck architecture, depicted in Figure 6, streamlines the feature fusion process between local and global models at the feature extraction stage. The input feature map has the dimensions C₁ × H × W. It is processed through the Conv module, resulting in an output feature map of C₁/2 × H × W. Subsequently, the GSconv module generates another output feature map of C₂/4 × H × W. The feature fusion takes place through the FCConv module, resulting in an output feature map of C₂/2 × H × W. In the final step, we concatenate the outputs of the Conv and FCConv modules to acquire an output feature map sized C₂ × H × W.

2.5. PSA Attention Mechanism

The attention mechanism known as the squeeze-and-excitation module (SE) [17] captures contextual information by Global Max Pooling (GMP) to determine the significance of different channels. However, the SE attention mechanism treats all neurons in a channel or spatial position equally, limiting its effectiveness in computing true 3D weights. Furthermore, the weight calculation in SE requires substantial computation, increasing the network’s computational cost and hindering its ability to extract more information from neurons.

To overcome these challenges, we introduce the parameter-free spatial attention mechanism (PSA) as a parameter-free alternative, as depicted in Figure 7. W and H, respectively, denote the feature map’s width and height. C represents the number of channels. The size of the input feature map is represented as W × H × C, where various colors correspond to distinct channels, spatial positions, or points on the feature map, each linked to an individual scalar value. The PSA distinguishes between different positions and treats all channels individually, assigning a unique weight to each neuron. Consequently, the PSA provides 3D weights that surpass the capabilities of traditional 1D and 2D attention mechanisms. Moreover, the PSA incorporates existing attention modules into each block, enhancing the output from the leading layer.

In order to incorporate both spatial and channel dimensions in 3D weights and effectively distinguish the target neuron from other, it is necessary to measure a value that distinguishes the target neuron from the rest, which enables the network to acquire more discriminative features. Consequently, we define the energy function for each neuron as follows:

$$e_t\left(w_t,b_t,y,x_i\right)=\frac{1}{N-1}{\displaystyle \sum _{i=1}^{N-1}{\left(-1-\left(w_t{x}_i+b_t\right)\right)}^2}+{\left(1-\left(w_{t}t+b_t\right)\right)}^2+\lambda w_t^2$$

(2)

where $w_t$ and $b_t$ are the weights and biases of the neuron transformation, $t$ is the target neuron, $x_i$ is the other neurons, i and j are spatial indices, and each channel has N = H × W energy functions. Therefore, the analytical solution for Equation (3) is:

$${}_{\begin{array}{l}{w}_t=-\frac{2\left(t-{\mu }_t\right)}{{\left(t-{\mu }_t\right)}^2+2{\sigma }_t^2+2\lambda }\\ b_t=-\frac{1}{2}\left(t+{\mu }_t\right)w_t\end{array}}$$

(3)

We can compute the minimum energy $e_t^{\ast }$ utilizing the following formula:

$$e_t^{\ast }=\frac{4\left(\widehat{\sigma }+\lambda \right)}{{\left(t-\widehat{\mu }\right)}^2+2{\widehat{\sigma }}^2+2\lambda }$$

(4)

Enhancement processing of features improve their quality or emphasize specific characteristics:

$$\tilde{X}=sigmoid\left(\frac{1}{E}\right)\otimes X$$

(5)

Among the steps in the enhancement processing, the operation E groups all $e_t^{*}$ values in both channel and spatial dimensions and applies a Sigmoid function to limit the value of E. Except for the calculation of channel mean $\mu$ and variance $\sigma$, all computations in this module are element-wise operations. This attention mechanism is unique compared to other attention mechanisms and greatly enhances the flexibility of the network.

2.6. Loss Function

The CIoU loss, as outlined in [18], can expedite the bounding box regression process. During the regression process, in the regression phase, if there exists a linear relationship between the aspect ratio of the predicted box and the ground truth box, the penalty term introduced in CIoU for the relative ratio loses its effectiveness. When one of the dimensions (either width or height) experiences an increase, the corresponding other dimension must decrease; it is not possible for both dimensions to increase or decrease simultaneously. To address this problem, we employ the EIoU loss function, which decomposes the loss into three components: the overlap loss L_EIoU measuring the overlap between the predicted box and the ground truth box, the center distance loss L_dis quantifying the distance between the centers of the predicted box and the ground truth box, and the aspect ratio loss L_asp characterizing the aspect ratio disparity between the predicted box and the ground truth box. The EIoU directly penalizes the predicted values of width w and height h, as shown below:

$$L_{\mathrm{EIoU}}=L_{\mathrm{IoU}}+L_{dis}+L_{asp}=1-\mathrm{IoU}+\frac{{\rho }^2\left(b,b^{gt}\right)}{c_w^2+c_h^2}+\frac{{\rho }^2\left(w,w^{gt}\right)}{c_w^2}+\frac{{\rho }^2\left(h,h^{gt}\right)}{c_h^2}$$

(6)

IoU denotes the intersection over union, while $c_w^2$ and $c_h^2$ represent the width and height of the minimum bounding rectangle for the predicted box and ground truth (GT) box, respectively. ${\rho }^2\left(b,b^{gt}\right)$ represents the Euclidean distance between the predicted box and the target box, ${\rho }^2\left(w,w^{gt}\right)$ denotes the difference in width between the predicted box and the target box, and ${\rho }^2\left(h,h^{gt}\right)$ represents the difference in height between the predicted box and the target box.

The first two components of the EIoU loss maintain the approach employed in the CIoU. However, the loss associated with width and height directly minimizes the dissimilarity between the predicted box and the ground truth box concerning their width and height, consequently accelerating the convergence process.

2.7. Evaluation Criteria

The evaluation metrics for IE-YOLOv5 include precision, recall, average precision (AP), mean average precision (mAP), and frames per second (FPS). FPS denotes the model’s processing capacity, indicating the number of images processed per second, which directly reflects the detection speed. Precision signifies the likelihood that all positive samples detected by the model are indeed positive samples, whereas recall quantifies the probability of detecting positive samples within the actual positive sample set. Equations (7) and (8) express precision and recall.

$$Precision=\frac{TP}{TP+FP}=\frac{TP}{predictions}$$

(7)

$$Recall=\frac{TP}{TP+FN}=\frac{TP}{groundtruths}$$

(8)

True positives (TP) denote the count of positive samples accurately predicted by the model and False positives (FP) signify the quantity of positive samples erroneously predicted by the model, while false negatives (FN) denote the quantity of negative samples overlooked by the model. After calculating the intersection over union (IoU) between each detection box and ground truth box, we can derive TP, FP, and FN values using a specified IoU threshold.

AP (average precision) denotes the area under the precision–recall curve pertaining to a particular class. It can be calculated using Equation (9), where K signifies the count of classes, P represents precision, and R represents recall. mAP (mean average precision) represents the mean value of AP across all classes, and it can be computed using Equation (10):

$$AP={\displaystyle {\int }_0^{1}P\left(R\right)dR}$$

(9)

$$mAP=\frac{{\displaystyle {\sum }_{i=1}^{K}A{P}_i}}{K}$$

(10)

3. Experiments

3.1. Experimental Data Set

The experimental dataset used is the NEU (Northeastern University) Steel Surface Defect Dataset [19]. This dataset consists of 6 categories of defect types, specifically crazing, inclusion, patches, pitted surface, rolled-in scale, and scratches, totaling 1800 defect images. Each class contains 300 images. To further enhance the NEU-DET dataset, we randomly selected 600 images from the training set. Due to the dataset having too few images and the images not being sufficiently distinct in terms of inter-class gaps as well as too large category gaps, we applied data augmentation techniques, including horizontal flipping, vertical flipping, and random cropping [20]. The dataset was expanded to 2000 images. Figure 8 displays the images for each defect class, and we divided the dataset into training, testing, and validation sets with a distribution ratio of 6:2:2.

3.2. Experimental Platform

All experiments were conducted on a unified platform as outlined in

Table 1. Training platform.

Experimental Equipment	Model Number
Appliances	Dell G15
Platforms	Pycharm 2020
CPU	Intel(R) Core (TM) i7-11800H
GPU	GeForce RTX 3060
System memory	16 GiB
SSD	Samsung SSD 970 EVO Plus 1 TB

.

3.3. Experimental Method

To assess the detection capabilities of the IE-YOLOv5 network using the NEU-DET defect detection dataset, ablation experiments were conducted to validate the optimization effects of each improvement module. Specifically, the following variants of YOLOv5 were compared:

• YOLOv5-A uses the lightweight convolution method GSConv to replace the standard convolution in the Neck module.
• YOLOv5-B builds upon YOLOv5-A by incorporating feature fusion using FF-Slim-Neck structure.
• YOLOv5-C implements the PSA attention mechanism at the conclusion of the backbone network.
• IE-YOLOv5 is the proposed model that combines all three improvement modules.

Through a comparative analysis of the performance of these various versions, we can evaluate the impact of each module on detection accuracy and efficiency.

To further substantiate the superiority of the proposed algorithm, we conducted comparative experiments with other widely recognized object detection algorithms, such as Faster R-CNN, SSD, YOLOv3, and YOLOv4. These algorithms find extensive application in the domain of object detection and serve as benchmarks for performance evaluation.

By comparing the detection accuracy, efficiency, and other metrics of the proposed algorithm (IE-YOLOv5) with these baseline algorithms, we can evaluate its performance and determine whether it outperforms or matches the most advanced methods in metal surface defect detection tasks.

4. Experimental Results and Analysis

4.1. Experimental Results

Based on the algorithm proposed in this study, the network model of IE-YOLOv5 was constructed. Figure 9a illustrates the training and validation loss curves of the IE-YOLOv5 network. The curves show a significant decrease in loss during the initial training stages. After approximately 100 iterations, both the training and validation loss curves stabilize and converge to around 0.2. A loss curve that stabilizes around 0.2 indicates that the model gradually converges to a relatively low level of loss during training. A stable loss curve means that the model is no longer improving significantly and that it has learnt the patterns in the training data.

Figure 9b illustrates the average accuracy value curve for the IE-YOLOv5 network. After approximately 80 iterations, both the training and validation sets exhibit a stable trend, with the final mean average precision exceeding 96%. An average accuracy of over 96% indicates that the model has performed very well on the test data and was able to accurately classify most of the samples. This indicates that the model has learnt the important features of the data and is able to achieve highly accurate predictions on new data.

Figure 10a depicts the comparison of mean average precision curves for YOLOv5, YOLOv5-A, YOLOv5-B, YOLOv5-C, and IE-YOLOv5. After approximately 120 iterations, YOLOv5 stabilizes, achieving a final mean average precision of around 94%. YOLOv5-A stabilizes after approximately 110 iterations, achieving a final mean average precision of around 95%. YOLOv5-B stabilizes after approximately 100 iterations, achieving a final mean average precision of around 96%. Similarly, YOLOv5-C stabilizes after approximately 100 iterations, achieving a final mean average precision of around 96%.

The IE-YOLOv5 network exhibits improved accuracy, with a respective increase of 2.4%, 1.2%, 0.6%, and 0.7%, compared to YOLOv5, YOLOv5-A, YOLOv5-B, and YOLOv5-C.

As illustrated in Figure 10b, it is evident that the enhanced model (IE-YOLOv5) achieves the convergence state more rapidly compared to the original YOLOv5 model as well as the YOLOv5-A, YOLOv5-B, and YOLOv5-C models. The original YOLOv5 algorithm exhibits significant instability in the initial stages, while the IE-YOLOv5 algorithm demonstrates smoother training loss curves. At the end of training, the training loss of IE-YOLOv5 was 27% lower compared to YOLOv5.

The findings demonstrate that the proposed approach minimizes loss and effectively addresses overfitting, resulting in a more stable and efficient training process for the IE-YOLOv5 model.

4.2. Analysis of Ablation Experiments

The model parameters of YOLOv5, YOLOv5-A, YOLOv5-B, YOLOv5-C, and IE-YOLOv5 are presented in

Table 2. Comparison of model parameters.

Model	GSconv	Slim-Neck	PSA	Weight	Inference Time
YOLOv5	✕	✕	✕	23.2 MB	2.8 ms
YOLOv5-A	✓	✕	✕	23.2 MB	2.7 ms
YOLOv5-B	✕	✓	✕	27.6 MB	2.5 ms
YOLOv5-C	✕	✕	✓	28.2 MB	2.7 ms
IE-YOLOv5	✓	✓	✓	28.2 MB	2.5 ms

✓ indicates that the module was introduced and ✕ indicates that it was not introduced.

. Table 2 presents the results, with the “Inference Time” column displaying the average time needed for the model to infer a single image, and the “Weight” column indicating the model’s volume or size.

Table 3 illustrates the results of the comparative experiments, showcasing how various model variants affect the performance of the YOLOv5 model:

• Model YOLOv5-A achieved notable improvements compared to the original YOLOv5 model. It showed a 3.6% reduction in inference time, a 0.8% increase in mean average precision (mAP), and a 1.6% boost in frames per second (FPS). This indicates that replacing the standard convolution (Conv) with GSConv not only reduced inference time but also improved the detection accuracy and speed.
• The YOLOv5-B model demonstrated a 1.3% mAP improvement over the original YOLOv5 model. It validates the effectiveness of incorporating the FF-Slim-Neck structure while keeping the model’s size and inference time nearly unchanged. The findings indicate that the FF-Slim-Neck structure enhanced the model’s ability to generalize and improved its accuracy in detecting defects.
• Model YOLOv5-C, which incorporated the PSA attention mechanism, showed significant improvements. It experienced a 21.6% increase in volume, which may have led to a 3.6% decrease in detection speed. However, this was balanced by a 1% increase in mAP, demonstrating the effectiveness and efficiency of the proposed PSA attention mechanism. The improved feature fusion network played a significant role in performance enhancement.

Table 3. Comparison of Evaluation Indicators of Models.

Model	AP						mAP	FPS
	Craze	Inclusion	Patch	Pit	Roll	Scratch
YOLOv5	93.9%	93.8%	94.7%	94.3%	93.4%	95.7%	94.3%	48.65
YOLOv5-A	95.2%	96.1%	95.2%	95.5%	94.2%	96.5%	95.5%	49.45
YOLOv5-B	95.7%	96.5%	95.6%	96.3%	95.1%	97.1%	96.1%	48.69
YOLOv5-C	94.6%	96.3%	95.3%	96.3%	95.6%	97.9%	96.0%	46.29
IE-YOLOv5	95.7%	97.2%	96.3%	96.5%	96.4%	98.1%	96.7%	46.17

Table 3. Comparison of Evaluation Indicators of Models.

Model	AP						mAP	FPS
	Craze	Inclusion	Patch	Pit	Roll	Scratch
YOLOv5	93.9%	93.8%	94.7%	94.3%	93.4%	95.7%	94.3%	48.65
YOLOv5-A	95.2%	96.1%	95.2%	95.5%	94.2%	96.5%	95.5%	49.45
YOLOv5-B	95.7%	96.5%	95.6%	96.3%	95.1%	97.1%	96.1%	48.69
YOLOv5-C	94.6%	96.3%	95.3%	96.3%	95.6%	97.9%	96.0%	46.29
IE-YOLOv5	95.7%	97.2%	96.3%	96.5%	96.4%	98.1%	96.7%	46.17

In this study, the comparison between the proposed IE-YOLOv5 model and the original YOLOv5 model yields the following results:

• The IE-YOLOv5 model shows a 21.6% increase in volume, which indicates more parameters and potentially better representation capacity.
• There was a 5.1% decrease in detection speed for the IE-YOLOv5 model compared with the original YOLOv5 model, suggesting a slight reduction in real-time inference performance.
• The inference time was significantly reduced by 10.7% compared to the original YOLOv5, which shows overall computational efficiency improvement.
• There was a 2.4% improvement in mean average precision (mAP) compared to original YOLOv5, which indicates enhanced detection accuracy across different defect types.
• The IE-YOLOv5 model demonstrated increased detection accuracy for inclusion, roll, and scratch defects by 3.4%, 3.0%, and 2.4%, compared to the original YOLOv5, respectively.

These findings demonstrate that the improved feature fusion network in the IE-YOLOv5 model effectively contributes to enhancing the algorithm’s performance, especially with regard to detection accuracy and inference efficiency.

The defect detection speed of IE-YOLOv5 exhibits a marginal decrement of 5.1% when compared to the original YOLOv5. This signifies that under identical hardware and environmental conditions, IE-YOLOv5 necessitates a slightly longer duration to process images and generate defect detection outcomes. This diminished detection speed could potentially impact real-time applications that require swift responses, such as defect detection tasks on production lines, where sluggish performance may lead to reduced production efficiency.

Despite the slight reduction in detection speed, IE-YOLOv5 showed a 10.7% reduction in inference time. This denotes that the computational time required for each model inference is truncated. A shortened inference time holds particular significance for applications that mandate frequent model inferences, as it can enhance the overall system efficiency.

IE-YOLOv5 has struck a balance between performance and efficiency, rendering it suitable for certain scenarios. However, for real-time applications necessitating high-speed detection, the adoption of swifter models or hardware acceleration may be imperative.

4.3. Algorithm Performance Analysis

Table 4. Detection results of different algorithms.

Model	AP						mAP	FPS
	Craze	Inclusion	Patch	Pit	Roll	Scratch
Faster R-CNN	92.1%	93.5%	90.4%	91.2%	93.6%	92.5%	92.2%	17.67
SSD	90.5%	92.1%	90.2%	93.5%	91.4%	90.7%	91.4%	28.89
YOLOV3	92.5%	93.9%	92.2%	93.8%	95.1%	94.8%	93.7%	34.18
YOLOV4	93.8%	94.7%	93.1%	96.3%	96.2%	97.1%	95.2%	40.19
IE-YOLOv5	95.7%	97.2%	96.3%	96.5%	96.4%	98.1%	96.7%	46.17

displays the detection outcomes obtained using different algorithms. As indicated by the results in Table 4, the proposed algorithm in this study demonstrates outstanding performance in detecting metal surface defects, achieving the best overall metrics. The average precision for three types of metal surface defects, namely inclusion, pit, and scratch, reaches an impressive 97%, surpassing SSD’s lower recall rate for dense defects and enhancing the detection accuracy by 5.3%.

Moreover, in terms of algorithm speed, the adoption of the slender-neck structure in this study leads to minimal computational overhead, resulting in an impressive detection frame rate of 46.17 fps. This is significantly higher compared to Faster R-CNN and SSD networks, with an improvement of 28.5 fps and 17.28 fps, respectively. The implementation of the slender-neck structure is a key contributing factor to this remarkable speed enhancement.

The application of the improved IE-YOLOv5 algorithm, trained using the proposed method, to NEU-DET defect detection is depicted in Figure 10. From Figure 11, the enhanced algorithm presented in this research is better suited for real-world production processes compared to the original network.

5. Application Notes

Successful implementation of the IE-YOLOv5 method in an industrial setting entails a series of crucial steps and considerations. The following are key points to take into account:

• Data collection and preparation: gather diverse image data representing industrial environments, including various types of metal defects and different environmental conditions. Cleaning, labelling, and pre-processing of data to prepare datasets for training and testing IE-YOLOv5 models.
• Model training: select IE-YOLOv5 as the base model and fine-tune it according to industrial requirements. Perform hyperparameter tuning to achieve optimal performance. Train the model on large-scale datasets to improve its generalization capabilities.
• Hardware and the infrastructure: choose the appropriate hardware platform, such as GPUs or edge devices, to support the real-time performance of IE-YOLOv5. Configure a suitable computer vision workstation or server for model training and inference.
• Model deployment: deploy the trained IE-YOLOv5 model onto target devices or systems in the industrial environment. Integrate the model with existing industrial automation systems to ensure smooth data flow and collaboration with other devices.
• Performance evaluation and fine-tuning: assess the performance of IE-YOLOv5 in the actual industrial environment, including accuracy, detection speed, and robustness. Fine-tune the model further based on specific task requirements.
• Security and privacy: address data security and privacy concerns to ensure compliance with regulations and standards for handling and storing image data. Implement measures to prevent potential attacks and data leaks.
• Operator training: provide appropriate training for personnel to ensure they can use the IE-YOLOv5 system correctly, understand its outputs, and take appropriate actions in response to detected defects.
• Maintenance and updates: establish a maintenance plan to regularly monitor and maintain the IE-YOLOv5 system to ensure stable performance and keep it up-to-date. Over time, update the model to adapt to new defect types or changes in environmental conditions.

When implementing the IE-YOLOv5 method, you may encounter the following challenges and considerations:

Data quality: Data quality is crucial for model training and deployment. Noise, inconsistency, or labeling errors can negatively impact model performance.

Real-time requirements: some industrial applications demand real-time performance, necessitating hardware and model capabilities to meet these requirements.

Hardware costs: selecting the right hardware can involve high costs, particularly for large-scale deployments.

Compliance: ensure that the implementation of IE-YOLOv5 complies with industrial and regulatory standards, especially in highly regulated industries such as healthcare and aerospace.

Environmental variability: industrial environments can be influenced by factors such as temperature, humidity, and lighting, which may adversely affect model performance.

False positives and negatives: models may produce false positives or negatives, potentially leading to issues in production.

Maintenance and updates: models require regular maintenance and updates to adapt to new defect types and changing environmental conditions.

Successful implementation of the IE-YOLOv5 method requires careful planning, resource allocation, and collaboration across cross-functional teams. Addressing these challenges and considering these factors will help ensure that defect detection systems in industrial environments operate efficiently and provide accurate results.

6. Conclusions

A novel method based on the Information Enhancement YOLOv5 Network (IE-YOLOv5) for detecting surface defects of metal part is proposed in this paper. A slender neck module, FF-Slim-Neck, which combined GSConv and federated learning-based feature fusion is introduced in IE-YOLOv5, which can realize a lightweight network with high detection efficiency. A parameter-free spatial attention module (PSA) was used in IE-YOLOv5, which can realize poor information feature extraction in crucial areas without the need for excessive structural adjustments. Overall optimization and adjustments were made to the backbone network of IE-YOLOv5, which significantly reduces the model’s parameters and improves its efficiency. Experiments were carried out with the NEU-DET public dataset from Northeastern University. The experimental results indicate that the proposed algorithm for detecting defects on metal surfaces achieves an average precision of 96.7% when identifying six different types of surface imperfection: crazing, inclusions, patches, pitting, scaling, and scratches. This represents a 2.4% enhancement in precision compared to the YOLOv5 algorithm. This enhancement aligns well with the demands of defect detection in industrial settings for metal parts. IE-YOLOv5 has important practical implications for the actual detection of metal defects on a number of levels, including increased productivity, reduced inspection times in the metal manufacturing process, which is essential for increased productivity, and more accurate detection and identification of metal defects resulting in a reduction in reject rates. Reducing defect rates not only saves money, but also improves product quality and improves safety, reducing the need for manual intervention and lowering the risk of operator injury.

In considering the future research directions for the IE-YOLOv5 method, we propose an exploration of semi-supervised or self-supervised learning approaches to reduce its dependency on large-scale annotated data. This endeavor aims to enhance the utility of IE-YOLOv5 in scenarios characterized by data scarcity. Additionally, we will investigate the fusion of multi-modal data from various sensors, such as images, audio, and temperature readings, into the IE-YOLOv5 framework to enhance defect detection and recognition capabilities. This approach holds the potential to elevate the algorithm’s robustness in complex environments.