Unleashing the power of AI in detecting metal surface defects: an optimized YOLOv7-tiny model approach

Depth Aware Convolution module

In the YOLOv7-tiny network, the ELAN-T module is used to extract features from feature maps of various regions. However, for our metal surface defect dataset, defects within the same category vary greatly in size and shape. Furthermore, there may also be similarities between defects of different categories. Moreover, due to the variations in different sample materials and the influence of lighting conditions, the grayscale values of intra-class defect images also undergo certain changes, which in turn disrupt the accuracy of the detection results. These factors collectively make it challenging for the network to extract meaningful features. The ELAN-T module is a simplified version of the ELAN module in YOLOv7. However, the ELAN module itself is not very effective in feature extraction (Wang et al., 2023b). Hence, it is clear that the feature extraction capability of ELAN-T is also unsatisfactory. Therefore, we attempt to improve the ELAN-T module of YOLOv7-tiny by introducing our innovative depth aware convolution module (DAC) (as shown in Fig. 3). We replace all ELAN-T modules in the entire model structure with DAC modules to further enhance the defect feature extraction and fusion capabilities in the YOLOv7-tiny network architecture while maintaining a certain inference speed.

The inspiration for improving the DAC module comes from the combination of V7 and V7-tiny themselves. Due to the simple structure of the ELAN-T (C5) module in V7-tiny, its feature extraction capabilities are significantly limited compared to other models in the V7 series; The ELAN-T module consists of five convolutions: three 1 × 1 convolutions and two 3 × 3 convolutions. It is known that 1 × 1 convolutions only facilitate information exchange and feature fusion between channels, lacking interactions and fusion between neighboring pixels. On the other hand, there are only two 3 × 3 convolutions that enhance feature fusion between neighboring pixels. Consequently, it is reasonable to expect a poor feature extraction capability from the ELAN-T module. In summary, we enhance the C5 module by increasing the number of 1 × 1 and 3 × 3 convolutions. The 1 × 1 convolutions strengthen the information exchange among channels, while the increased 3 × 3 convolutions widen the receptive field, promoting better interaction of spatial information and feature fusion. Additionally, we employ the Concat operation to increase the number of features, consequently improving the non-linear transformation of the network.

AWFP-add module

YOLOv7-tiny incorporates the feature fusion network of the YOLOv5 series, which consists of the Feature Pyramid Network (FPN) and the Path Aggregation Network (PAN) architecture (Wang, Bochkovskiy & Liao, 2023a). Lin et al. (2017) effectively transfers strong semantic information from deeper feature layers to even deeper layers (Chen et al., 2021). On the other hand, PAN transmits accurate localization information from bottom to top. By combining FPN and PAN, different detection layers from the backbone are parameterized together, enhancing the feature fusion capability of the network. However, this combination introduces a drawback: the PAN structure takes as input the features that have been processed by the FPN, but some defect features extracted from the backbone’s original information are lost along the way. The lack of original information for learning can lead to bias in training and consequently affect detection accuracy. To address this issue, we propose incorporating the concepts of bi-directional weighted feature pyramid network (BiFPN) and fast normalization fusion from its associated paper into the conventional Add module. This allows the Add module to have its own learnable parameters. Additionally, in terms of structure, we insert the Add module after each of the three Concat modules in the feature fusion network. We also combine the practical significance of surface defect detection in this study and the concept of residuals. Specifically, we connect the Add module and Backbone to the three feature maps provided to the Head, enabling the network to retain more shallow semantic information without losing too much deep semantic information. At the same time, different weights are set according to the importance of different input features to ensure that the network can adjust the weight parameters adaptively during gradient backhaul, and pay attention to the importance of different features and adjust them. This innovation is the second improvement described in this article, AWFP-Add, and the formula and structure are shown in Fig. 4:

(1)${\displaystyle O=\sum _i\frac{w_i}{\varepsilon +\sum _j{w}_j}\cdot I_i}$ (2)${\displaystyle k_i=Sigmoid\left(w_i\right)}$ (3)${\displaystyle Z=\left(1-k\right)\ast y+k\ast x.}$

Equation (1) is the mathematical expression of fast normalized fusion, for this study, directly introducing fast normalized fusion can improve model accuracy to a certain extent, but the improvement is very limited. Moreover, it introduces more learnable parameters, not only increasing the model’s computational complexity but also significantly affecting the model’s stability. For this reason, we opt to adopt the simple and efficient sigmoid function to restrict the weight value range between (0, 1) and implement the task of the network to focus on the importance of different features by itself through Eq. (3), and its complete computational process is illustrated in Fig. 5.

Focal-SIoU

The loss function of YOLOv7-tiny consists of three components: bounding box regression loss (BBR), confidence loss function, and class loss function. The bounding box regression loss measures the error in the predicted box’s coordinate localization error. The confidence loss reflects the confidence error of the predicted box, while the class loss function captures the prediction box’s error in predicting the target class. Specifically, the class loss function uses binary cross-entropy (BCE) loss, only calculating the classification loss of positive samples. The confidence loss function is also BCE loss, but it measures the confidence loss between the predicted bounding box and the ground truth box using complete intersection over union (CIOU). This loss is calculated for all samples. The bounding box loss function also uses CIOU loss, but it only calculates the position loss of positive samples. The CIOU loss takes into account three important geometric factors: overlap area, the distance between centroid points, and aspect ratio, which makes the bounding box regression more stable. This is shown in Fig. 6. Given the predicted box B and the ground truth box B_gt. The CIOU loss is defined by Eq. (4): (4)${\displaystyle {Loss}_{\text{CIOU}}=1-\mathrm{IOU}+\frac{{\rho }^2\left(b,b^{gt}\right)}{c^2}+\alpha v}$

where ${\rho }^2\left(b,b^{gt}\right)$ denotes the Euclidean distance between the center point of the prediction box and the center point of the ground truth box, denoted by d in Fig. 6, c denotes the diagonal distance between the ground truth box and the smallest closed rectangle contained in the prediction box. α is defined by Eq. (5) below and ν is defined by Eq. (6) below.

(5)${\displaystyle \alpha =\frac{v}{1-\mathrm{IOU}+v}}$ (6)${\displaystyle v=\frac{4}{{\pi }^2}{\left(arctan\frac{w^{gt}}{h^{gt}}-arctan\frac{w}{h}\right)}^2}$

where w_gt represents the width of the ground truth box, h_gt represents the height of the ground truth box, w represents the width of the prediction box, h represents the height of the prediction box. However, there is still room for improvement in the boundary box loss function. For instance, the CIOU does not take into account the orientation between the ground truth and predicted bounding boxes, which results in slower convergence. To address this, we propose replacing the original CIOU loss function with scylla intersection over union (SIOU), which introduces the vector angle between the ground truth and predicted bounding boxes to redefine correlation (Gevorgyan, 2022). The SIOU loss function consists of four components: distance loss, angle loss, shape loss, and IOU loss.

(1) ANGLE COST

The angle cost is defined by Eq. (7). Its diagram is shown in Fig. 7. (7)${\displaystyle \Lambda =1-2\times {sin}^2\left(arcsin\left(\frac{c_h}{\sigma }\right)-\frac{\pi }{4}\right)}$

where c_h is the height distance between the center point of the ground truth box and the prediction box, σ is the distance of the center point between the ground truth box and the prediction box.

(8)${\displaystyle c_h=max\left(b_{c_y}^{gt},b_{c_y}\right)-min\left(b_{c_y}^{gt},b_{c_y}\right)}$ (9)${\displaystyle \sigma =\sqrt{{\left(b_{c_x}^{gt}-b_{c_x}\right)}^2+{\left(b_{c_y}^{gt}-b_{c_y}\right)}^2}}$

where $\left(b_{c_x}^{gt},b_{c_y}^{gt}\right)$ is the center coordinate of the ground truth box and $\left(b_{c_x},b_{c_y}\right)$ is the center coordinate of the prediction box.

(2) DISTANCE COST

The distance cost is defined by Eq. (10).Its diagram is shown in Fig. 8.

(10)${\displaystyle \Delta =\sum _{t=x,y}\left(1-e^{-\gamma p_t}\right)=2-e^{-\gamma {\rho }_x}-e^{-\gamma p_y}}$ (11)${\displaystyle {\rho }_x={\left(\frac{b_{c_x}^{gt}-b_{c_x}}{c_w}\right)}^2,{\rho }_y={\left(\frac{b_{c_y}^{gt}-b_{c_y}}{c_h}\right)}^2}$ (12)${\displaystyle \gamma =2-\Lambda }$

where (cw, ch) is the width and height of the minimum external matrix of the ground truth box and the prediction box.

(3) SHAPE COST

The shape cost is defined by Eq. (13).

(13)${\displaystyle \Omega =\sum _{t=w_{2}h}{\left(1-e^{-w_t}\right)}^{\theta }={\left(1-e^{-w_w}\right)}^{\theta }+{\left(1-e^{-w_h}\right)}^{\theta }}$ (14)${\displaystyle w_w=\frac{\left|w-w^{gt}\right|}{max\left(w,w^{gt}\right)},w_h=\frac{\left|h-h^{gt}\right|}{max\left(h,h^{gt}\right)}}$

where (w, h) is the width and height of the prediction box, $\left(w^{gt},h^{gt}\right)$ is the width and height of the ground truth box, θ controls the degree of attention to shape loss.

(4) IoU COST

The IoU cost is defined by Eq. (15). (15)${\displaystyle \mathrm{IoU}=\frac{\mathrm{A}}{\mathrm{B}}}$

where A represents the intersection of the ground truth box and the prediction box, B represents the union of the ground truth box and the prediction box.

(5) SIoU LOSS

In conclusion, the SIoU loss is defined by Eq. (16). (16)${\displaystyle {Loss}_{\text{SIOU}}=1-\mathrm{IOU}+\frac{\Delta +\Omega }{2}.}$

In boundary box regression, the problem of imbalanced training samples also arises, where the sparsity of the target objects in the images leads to a scarcity of high-quality examples with small regression errors compared to low-quality examples. To concentrate the SIOU loss on high-quality examples, this study considers combining the focal loss, which specifically deals with the imbalance between positive and negative samples, with the SIOU loss (Zhang et al., 2022). Therefore, we propose the Focal-SIOU loss to enhance the performance of the SIOU loss, and it is defined by Eq. (17). (17)${\displaystyle L_{\text{Focal-SIOU}}={\text{IOU}}^{\gamma }{L}_{\text{SIOU}}.}$

Experiments on GC10-DET

From Table 5, it is evident that our improvements have been effective. The original YOLOv7-tiny model (baseline) achieved an mAP of 70.2. In comparison, C-YOLO outperforms it with a 4.3 higher mAP. Furthermore, C-YOLO demonstrates superior detection accuracy than YOLOv7-tiny in all defect categories except for Pu. CF-YOLO and CA-YOLO achieve mAPs of 75.3 and 77.1 respectively, surpassing the original model by 5.1 mAP and 6.9 mAP. CAF-YOLO attains an mAP of 81%, the highest among them. Regarding specific defect types, CAF-YOLO boasts the highest performance in all defect detections except for Wl. The comparison of our improvement with the baseline model on the PR curve is shown in Fig. 14, confirming that our modification can identify various defects. CF-YOLO increases the mAP by +4.3 to +5.1, indicating that the mode of Focal Loss+SIoU to some extent mitigates the issue of imbalance between positive and negative examples in bounding box regression, as well as inaccurate regression boxes. By introducing learnable parameters in the Add module and incorporating the architecture of adaptive weighted fusion paths, the improvement range expands from +4.3 to +6.9, reflecting the positive impact of the AWFP-Add module.

Experiments on NEU-DET

To further investigate the effectiveness and robustness of our proposed method, we conducted ablation experiments on NEU-DET. Table 6 presents the results, showing that the baseline model YOLOv7-tiny achieved an mAP of 70.9. Meanwhile, C-YOLO, CF-YOLO, and CA-YOLO achieved mAP scores of 74.2, 76.3, and 76.7, respectively. The highest mAP is implemented by CAF-YOLO, and its comparison with the base model on the PR curve, as shown in Fig. 15, most accurately identifies defects such as Patches, Scratches, etc. The results indicate that when we use the DAC module, the accuracy of most types of defect detection increases. Nevertheless, the AWFP-Add module achieved a 2.5 mAP boost, and the use of Focal-SIoU resulted in a performance improvement of 2.1 mAP, which fully demonstrates the effectiveness of our modified Backbone and Head components.