Improved YOLOv8-based truck floor damage detection method
By improving the YOLOv8 algorithm and a multi-model series detection method combined with ResNet classification model, the problem of high cost and slow speed of railway truck floor damage detection in the prior art is solved, and fast, accurate and efficient fault detection is achieved.
Patent Information
- Application Number
- CN202411886264.1
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-12-20
- Publication Date
- 2025-04-11
AI Technical Summary
Existing deep learning methods are used for railway truck floor damage detection, which have high cost, slow detection speed and low efficiency, making it difficult to meet the needs of fast and accurate fault identification.
The truck floor damage detection method based on improved YOLOv8 is adopted, and the segmented data set is generated through data augmentation and data annotation. The YOLOv8 algorithm is optimized by PConv convolution, and multi-model series detection is combined with the ResNet classification model to improve detection accuracy and efficiency.
It significantly improves detection speed, robustness and generalization, reduces detection costs, and achieves high accuracy and high efficiency fault detection to meet the actual needs of the railway transportation industry.
Smart Images
Figure CN119810415A_ABST
Abstract
Description
Technical Field
[0001] The present invention belongs to the field of damage detection, and in particular relates to a truck floor damage detection method based on improved YOLOv8. Background Art
[0002] The damaged floor of railway freight cars is a common fault that poses a serious threat to driving safety. Traditional manual detection methods have been unable to meet the needs of quickly and accurately identifying faults. In recent years, the development of deep learning and artificial intelligence technology has provided new possibilities for solving this problem. Using deep learning algorithms, railway freight car floors can be automatically inspected, thereby reducing human misjudgment and missed detection, and improving the accuracy of fault detection.
[0003] However, current deep learning methods also have some problems, such as the need for expensive GPU equipment investment, slow detection speed, and low efficiency. Therefore, it is particularly urgent and necessary to develop a new method that can both ensure fault detection accuracy and reduce detection costs. Summary of the invention
[0004] The present invention provides a truck floor damage detection method based on improved YOLOv8, which can effectively solve the problems in the background technology.
[0005] The present invention provides a truck floor damage detection method based on improved YOLOv8, comprising the following steps:
[0006] A1: Collect full-vehicle linear array images and fault sample images of the bottom of different models of railway freight cars under different conditions;
[0007] A2: Perform data enhancement and data annotation on the collected fault sample images to obtain a segmented data set;
[0008] A3: Optimize and improve the original YOLOv8 algorithm based on PConv convolution, and use the segmentation dataset in step A2 to train the improved YOLOv8 segmentation model;
[0009] A4: Use the trained improved YOLOv8 segmentation model to perform floor segmentation and fault identification on the image;
[0010] A5: Use the floor segmentation results to filter out the fault identification results of the non-floor area, and judge the credibility of the fault area according to the confidence score. If the credibility is met, the result is reported; if the credibility is insufficient, proceed to step A6;
[0011] A6: Use the trained improved YOLOv8 segmentation model to segment and detect the full-vehicle linear array image of the truck bottom in step A1, and crop all local images of possible floor damage areas and save them as a classification data set;
[0012] A7: Use the constructed classification dataset to train the ResNet classification model;
[0013] A8: The damaged area of the chassis of the image that does not meet the credibility in step A5 is cropped, scaled to a standard size, and then sent to the trained ResNet classification model for detection;
[0014] A9: The test results that meet the credibility requirements will be reported; otherwise, they will be considered as false detections and will not be reported.
[0015] As a further optimization of the present invention, the improved YOLOv8 algorithm of step three includes the following steps:
[0016] B1: The training image is input into the backbone network and processed by the stage1, stage2, stage3, stage4, stage5 and SPPF modules in sequence, and the feature map F3, feature map F4 and feature map F5 are output by the PC2 module in the stage3, the PC2 module in the stage4 and the SPPF module respectively;
[0017] B2: The feature maps F3, F4 and F5 extracted by the backbone network are sent to the neck part of the PAN-FPN structure, and the low-level physical features and high-level semantic features are aggregated at multiple scales using bottom-up path connections, thereby constructing a pyramid feature map and providing output feature maps C3, C4 and C5 for the detection head;
[0018] B3: Send feature map C3 to the detection prediction head of YOLO, and output the training detection result of the first level size after processing; send feature map C4 to the detection prediction head of YOLO, and output the training detection result of the second level size after processing; send feature map C5 to the detection prediction head of YOLO, and output the training detection result of the third level size after processing;
[0019] B4: Use the training set data and the set configuration to train the model, and calculate the loss function of the detection results and the target's true label on the validation set data. Use the SGD learning optimizer to update the network weights. After multiple iterative training, obtain the model weight file that can accurately segment the target.
[0020] As a further optimization of the present invention, step B2 includes the following method:
[0021] The feature map F5 is upsampled and then concatenated with the feature map F4 to form a feature map P4; the feature map P4 is subjected to a feature extraction module PC2 to obtain a feature map P4'; the feature map P4' is parallel upsampled and then concatenated with the feature map F3 to form a feature map P3; the feature map P3 is outputted as a feature map C3 by the feature extraction module PC2;
[0022] The output feature map C3 is passed through the downsampling PConv convolution module and then concatenated with the feature map P4', and then passed through the feature extraction module PC2 to output the feature map C4;
[0023] The output feature map C4 is passed through the downsampling PConv convolution module and then concatenated with the feature map F5, and then passed through the feature extraction module PC2 to output the feature map C5.
[0024] According to the method for detecting truck floor damage based on improved YOLOv8 in claim 2, it is characterized in that the downsampling module PConv in each stage in step B1 reduces the width and height of the input feature map to half of the original, and increases the number of feature channels to twice the original, and then the feature fusion module PC2 performs feature extraction.
[0025] As a further optimization of the present invention, when performing feature extraction, for an input feature map of size [h, w, c], the first basic PConv module performs a dimensionality reduction operation on the input feature map so that the channel dimension c is half of the original; it is then processed by two stacked P-Bottlenecks, and the output features are residually spliced with the original features to restore the dimension to c; then the SimAM attention mechanism module extracts key information from the spliced features, and the last basic PConv module adjusts the features to a fixed output feature dimension to complete the feature output.
[0026] As a further optimization of the present invention, the P-bottleneck module is composed of two stacked basic PConv modules and residual connections, the convolution kernel size of each PConv module is set to 3, the stride is 1, and the residual connection used is an addition operation.
[0027] As a further optimization of the present invention, a SimAM attention mechanism module is embedded after the P-Bottleneck module, and the 3-D attention weight of the feature map is inferred in a parallel branch, a unique weight is assigned to each neuron, and the minimum energy function is derived to explore the importance of each neuron; the expression formula for enhancing the feature is:
[0028]
[0029] In the formula, X is the concatenated input feature map; sigmoid is the activation function to limit the maximum value of E; E is the value of all neurons in the attention feature map A collection of; represents the minimum energy, and its expression is:
[0030]
[0031] in, λ is a parameter and t is a neuron.
[0032] The present invention provides a truck floor damage detection method based on improved YOLOv8, which has the following beneficial effects:
[0033] 1. Use lightweight model structure and attention mechanism structure to build neural network, which significantly improves the detection speed, robustness and generalization of the model;
[0034] 2. The multi-model serial detection method has higher detection accuracy and can effectively reduce misjudgment and missed detection;
[0035] 3. The model has a small number of parameters, which is suitable for low-cost equipment deployment and reduces investment costs;
[0036] 4. It can realize automated fault detection with high detection recognition rate and fast recognition speed, meeting the actual needs of the railway transportation industry and improving the safety and transportation efficiency of railway freight.
[0037] Instruction Manual
[0038] Figure 1 is a schematic diagram of the process of this embodiment;
[0039] Figure 2 This is a schematic diagram of the improved YOLOv8 segmentation neural network model of this embodiment;
[0040] Figure 3 This is a structural diagram of the Partial Convolution module of this embodiment;
[0041] Figure 4 This is the SimAM module structure diagram of this embodiment. DETAILED DESCRIPTION
[0042] like Figure 1 As shown, the present invention provides a truck floor damage detection method based on improved YOLOv8, comprising the following steps:
[0043] A1: Collect full-vehicle linear array images and fault sample images of the bottom of railway freight cars of different models under different conditions.
[0044] Grayscale images of the bottom of a moving railway freight car are obtained using high-definition linear array cameras mounted around the freight car tracks. Since the floors of freight cars of different models are different, and the quality of images taken by the equipment may be affected by natural weather conditions such as light and rain, in order to ensure the robustness of the training model, it is necessary to collect images of the bottom floor of the whole vehicle of different models under various weather conditions such as sunny days, rainy days, and cloudy days to ensure the diversity of training data. In addition, it is necessary to collect images of floor damage faults from the long-term collected images of the bottom floor of the whole vehicle, as well as images of water stains, oil stains, mud stains, etc. that are easily confused with floor damage faults as sample images for training and optimizing the recognition algorithm to improve its accuracy and robustness in practical applications.
[0045] A2: Perform data enhancement and data annotation on the collected fault sample images to obtain a segmented data set.
[0046] Manually collected data is not only costly, but also often difficult to cover all possible situations. These data may still have certain limitations. In order to overcome these limitations, data enhancement operations are performed on the collected grayscale images, including but not limited to small-angle image rotation, horizontal flipping, scaling, histogram equalization, random noise, brightness adjustment and other methods. In specific implementation, a random one or a combination of several methods will be performed on the grayscale image to increase the number and diversity of labeled samples and improve the model's generalization ability for different scenes, angles and scales.
[0047] Next, the Labelme tool is used to annotate all the processed sample images. The annotation categories are divided into three categories: floor area, damage and failure, and other categories (negative samples such as water stains, oil stains, and mud stains). The annotation results are mask images corresponding to each category.
[0048] The labeled sample images are divided into training set, validation set and test set in a ratio of 8:1:1 for training the neural network model.
[0049] A3: The original YOLOv8 algorithm is optimized and improved based on PConv convolution, and the neural network model is trained using the segmented dataset.
[0050] The original images in the prepared segmentation dataset are adaptively scaled to 640×640×3 and sent to the improved YOLOv8 neural network model for training. The prediction loss is calculated based on the training test results and the target true label, and the segmentation model parameters are optimized using the prediction loss. The trained semantic segmentation neural network model is output after repeated iterations until the best accuracy is achieved on the test set.
[0051] The improved YOLOv8 model includes using lightweight PConv convolution operations and SimAM parameter-free attention mechanism to reconstruct the backbone network and neck part of the model, as well as the YOLO detection head for final prediction. The improved model framework is as follows: Figure 2 The improved model architecture uses PConv convolution for feature extraction, which significantly reduces the number of parameters and computations of the entire network. In addition, a parameter-free SimAM attention mechanism module is embedded in each PC2 feature extraction module in the backbone network, which efficiently condenses image features without increasing network complexity, making feature fusion simpler and more efficient, improving detection speed and reducing model complexity.
[0052] The training specifically includes the following steps:
[0053] B1: The backbone network of the model is reconstructed using lightweight PConv convolution operations and SimAM parameter-free attention mechanism.
[0054] The backbone network includes five feature extraction stages with the same configuration, stage1, stage2, stage3, stage4, stage5, and an SPPF attention module. The feature extraction stage stage includes a downsampling module PConv and a feature fusion module PC2.
[0055] The training image is input into the backbone network and processed by stage1, stage2, stage3, stage4, stage5 and SPPF modules in sequence; the downsampling PConv module in each stage will reduce the width and height of the input feature map to half of the original, and increase the number of feature channels to twice the original, and then the feature fusion module PC2 will perform feature extraction; then the feature map F3, feature map F4 and feature map F5 are output by the PC2 module in the stage3 stage, the PC2 module in the stage4 stage and the SPPF module.
[0056] In this embodiment, the PConv convolution module is composed of a 2D PConv convolution, a BN layer, and a SiLU activation function connected in series to form a basic convolution module of the improved model. The 2D PConv convolution operation execution process is as follows: Figure 3 When performing downsampling convolution, the convolution kernel size of the 2D PConv convolution is set to 3 and the stride is set to 2.
[0057] The feature fusion module PC2 is composed of two basic PConv modules, two P-Bottleneck modules, residual connection and SimAM attention mechanism modules. The specific structure is as follows Figure 2 as shown in .
[0058] When performing feature extraction, for an input feature map of size [h, w, c], the first basic PConv module performs a dimensionality reduction operation on the input feature map so that the channel dimension c is half of the original; it is then processed by two stacked P-Bottleneck modules, and the output features are residually concatenated with the original features to restore the dimension to c; then the SimAM attention mechanism module extracts key information from the concatenated features, and the last basic PConv module adjusts the features to a fixed output feature dimension to complete the feature output.
[0059] In this embodiment, the P-bottleneck module consists of two stacked basic PConv modules and residual connections. The convolution kernel size of each PConv module is set to 3, the stride is 1, and the residual connection used is an addition operation instead of a concatenation operation to ensure that the dimension of the feature remains unchanged.
[0060] Specifically, the SimAM attention mechanism module is embedded after the P-Bottleneck module to help each feature extraction module of the backbone network extract feature information more accurately and increase the interpretability of the model. Figure 4 As shown in the figure, this operation does not need to add parameters to the original network, but instead infers the 3D attention weights of the feature map in a parallel branch, assigns a unique weight to each neuron, and derives the minimum energy function to explore the importance of each neuron; the expression formula for enhancing the feature is:
[0061]
[0062] In the formula, X is the concatenated input feature map; sigmoid is the activation function to limit the maximum value of E; E is the value of all neurons in the attention feature map A collection of; represents the minimum energy, and its expression is:
[0063]
[0064] in, The lower the energy, the greater the difference between neuron t and the surrounding neurons, the higher the importance, and the corresponding assigned weight should be higher.
[0065] In this embodiment, for the original image with an input size of 640×640×3, the size of the feature map F3 output by the backbone network is [80, 80, 128], the size of the feature map F4 is [40, 40, 256], and the size of the feature map F5 is [20, 20, 512].
[0066] B2: The feature maps F3, F4 and F5 extracted by the backbone network are sent to the neck part of the FPN+PAN structure, and the low-level physical features and high-level semantic features are aggregated at multiple scales using bottom-up path connections to construct a pyramid feature map and provide output feature maps C3, C4 and C5 for the detection head.
[0067] The specific method is as follows:
[0068] After upsampling the feature map F5, it is concatenated with the feature map F4 to form the feature map P4; the feature map P4 is subjected to the feature extraction module PC2 to obtain the feature map P4'; after upsampling the feature map P4' in parallel, it is concatenated with the feature map F3 to form the feature map P3; the feature extraction module PC2 is used to output the feature map C3 for the feature map P3;
[0069] The output feature map C3 is passed through the downsampling PConv convolution module and then concatenated with the feature map P4', and then passed through the feature extraction module PC2 to output the feature map C4;
[0070] The output feature map C4 is passed through the downsampling PConv convolution module and then concatenated with the feature map F5, and then passed through the feature extraction module PC2 to output the feature map C5.
[0071] In this embodiment, the feature extraction module PC2 used in the Neck part is no longer embedded in the attention mechanism module SimAM.
[0072] In this embodiment, the size of the output feature map C3 is [80, 80, 128], the size of the output feature map C4 is [40, 40, 256], and the size of the output feature map C5 is [20, 20, 512].
[0073] B4: Send feature map C3 to the detection prediction head of YOLO, and output the training detection result of the first level size after processing; send feature map C4 to the detection prediction head of YOLO, and output the training detection result of the second level size after processing; send feature map C5 to the detection prediction head of YOLO, and output the training detection result of the third level size after processing.
[0074] B5: Use the training set data and the set configuration to train the model, and calculate the loss function of the detection results and the target's true label on the validation set data. Use the SGD learning optimizer to update the network weights. After multiple iterations of training, obtain the model weight file that can accurately segment the target.
[0075] Among them, the loss function consists of classification loss function, positioning loss function and segmentation mask loss function; specifically, the classification loss and segmentation loss use the cross entropy loss function (BCE Loss), and the positioning loss uses DFL Loss and CIOULoss loss functions.
[0076] A4: Use the improved YOLOv8 segmentation model to perform floor segmentation and fault identification on the image.
[0077] The trained improved YOLOv8 segmentation model is used to perform fault detection on the floor image to be detected. The segmentation results of the floor area, the damaged floor area, and other categories of areas can be obtained, and the binary segmentation image can be further obtained based on the segmentation results.
[0078] A5: Use the floor segmentation results to filter out the fault identification results of the non-floor area, and judge the credibility of the fault area based on the confidence score. The fault area with insufficient credibility is detected again using the two-stage classification model.
[0079] The binary segmentation images of the floor area and the floor damage area are combined to determine whether the floor damage is within the floor area. If it is not within the floor area, the fault detection is considered to be a false detection and the detection result is discarded; other categories of areas detected are also considered to be false detections and the results are not reported.
[0080] After the first round of judgment, the floor damage fault detection result in the floor area is still retained, and a second round of judgment is further performed based on the detection confidence score: if the confidence score is greater than 0.8, it is considered to be a fault and the result is reported directly; if the confidence score is lower than 0.8, it indicates that the model cannot determine whether it is a fault, which requires secondary detection.
[0081] The situations that cause misdetection of damaged floor faults in railway freight cars are very complex. Although the improved YOLOv8 segmentation model has been trained with some negative sample data such as water stains, oil stains, and mud, this cannot cover all situations. Therefore, there are still misdetections in the damaged floor fault areas with confidence scores lower than 0.8. In order to avoid modifying the trained segmentation model, a classification model is used for further detection of such situations.
[0082] A6: Use the trained improved YOLOv8 segmentation model to segment and detect the full-vehicle linear array image of the bottom of the truck in step A1, and crop and save the local images of all possible floor damage areas as the classification data set.
[0083] The improved YOLOv8 segmentation model was used to segment and detect the original floor images of a large number of railway freight cars, and all local images of possible floor damage areas were cropped and saved as classification data sets. These classified images were manually checked and divided into two categories: floor damage and background.
[0084] A7: Use the constructed classification dataset to train the ResNet classification model.
[0085] A8: The damaged area of the chassis of the image that does not meet the credibility in step A5 is cropped, scaled to a standard size, and then sent to the trained ResNet classification model for detection.
[0086] Since the ResNet classification model performs secondary detection and judgment on the detection results of the segmentation network, the classification data set should also come from the detection results of the segmentation network.
[0087] ResNet-18 was selected as the network backbone, and the training size of the image was set to 224x224. Since the input image size of the classification model is small and the model itself is very small, the speed meets the real-time requirements, which is very suitable for the detection of multiple floor damage.
[0088] A9: The test results that meet the credibility requirements will be reported; otherwise, they will be considered as false detections and will not be reported.
[0089] like Figure 2 As shown in the figure, the specific process of using the ResNet classification model to identify floor damage faults is as follows:
[0090] The damaged floor segmentation area with a confidence score lower than 0.8 in the detection image is cropped, the image is scaled to a standard 224x224 size, and sent to the classification model for classification detection to obtain the classification result of the area, including the detection category and classification confidence score. If the category is damaged floor and the classification confidence score is greater than 0.8, it is considered a fault and the result is reported; otherwise, it is considered a false detection and the result is not reported.
[0091] Finally, it should be noted that the above embodiments are only used to illustrate the technical solution of the present invention, rather than to limit the scope of protection of the present invention. Although the present invention has been described in detail with reference to the preferred embodiments, those skilled in the art should understand that the technical solution of the present invention can be modified or replaced by equivalents without departing from the essence and scope of the technical solution of the present invention.
Claims
1. A truck floor damage detection method based on improved YOLOv8, characterized in that: The following steps are involved: A1: Collect full-vehicle linear array images and fault sample images of the bottom of different models of railway freight cars under different conditions; A2: Perform data enhancement and data annotation on the collected fault sample images to obtain a segmented data set; A3: Optimize and improve the original YOLOv8 algorithm based on PConv convolution, and use the segmentation dataset in step A2 to train the improved YOLOv8 segmentation model; A4: Use the trained improved YOLOv8 segmentation model to perform floor segmentation and fault identification on the image; A5: Use the floor segmentation results to filter out the fault identification results of the non-floor area, and judge the credibility of the fault area according to the confidence score. If the credibility is met, the result is reported; if the credibility is insufficient, proceed to step A6; A6: Use the trained improved YOLOv8 segmentation model to segment and detect the full-vehicle linear array image of the truck bottom in step A1, and crop all local images of possible floor damage areas and save them as a classification data set; A7: Use the constructed classification dataset to train the ResNet classification model; A8: The damaged area of the chassis of the image that does not meet the credibility in step A5 is cropped, scaled to a standard size, and then sent to the trained ResNet classification model for detection; A9: The test results that meet the credibility requirements will be reported; otherwise, they will be considered as false detections and will not be reported.
2. According to claim 1, a truck floor damage detection method based on improved YOLOv8 is characterized in that: Step 3 of the improved YOLOv8 algorithm includes the following steps: B1: The training image is input into the backbone network and processed by the stage1, stage2, stage3, stage4, stage5 and SPPF modules in sequence, and the feature map F3, feature map F4 and feature map F5 are output by the PC2 module in the stage3, the PC2 module in the stage4 and the SPPF module respectively; B2: The feature maps F3, F4 and F5 extracted by the backbone network are sent to the neck part of the PAN-FPN structure, and the low-level physical features and high-level semantic features are aggregated at multiple scales using bottom-up path connections, thereby constructing a pyramid feature map and providing output feature maps C3, C4 and C5 for the detection head; B3: Send feature map C3 to the detection prediction head of YOLO, and output the training detection result of the first level size after processing; send feature map C4 to the detection prediction head of YOLO, and output the training detection result of the second level size after processing; send feature map C5 to the detection prediction head of YOLO, and output the training detection result of the third level size after processing; B4: Use the training set data and the set configuration to train the model, and calculate the loss function of the detection results and the target's true label on the validation set data. Use the SGD learning optimizer to update the network weights. After multiple iterative training, obtain the model weight file that can accurately segment the target.
3. According to claim 2, a truck floor damage detection method based on improved YOLOv8 is characterized in that: Step B2 includes the following methods: The feature map F5 is upsampled and then concatenated with the feature map F4 to form a feature map P4; the feature map P4 is subjected to a feature extraction module PC2 to obtain a feature map P4'; the feature map P4' is parallel upsampled and then concatenated with the feature map F3 to form a feature map P3; the feature map P3 is outputted as a feature map C3 by the feature extraction module PC2; The output feature map C3 is passed through the downsampling PConv convolution module and then concatenated with the feature map P4', and then passed through the feature extraction module PC2 to output the feature map C4; The output feature map C4 is passed through the downsampling PConv convolution module and then concatenated with the feature map F5, and then passed through the feature extraction module PC2 to output the feature map C5.
4. According to claim 2, a truck floor damage detection method based on improved YOLOv8 is characterized in that: The downsampling module PConv in each stage of step B1 will reduce the width and height of the input feature map to half of the original, and increase the number of feature channels to twice the original, and then the feature fusion module PC2 will perform feature extraction.
5. According to claim 2, a truck floor damage detection method based on improved YOLOv8 is characterized in that: When performing feature extraction, for an input feature map of size [h, w, c], the first basic PConv module performs a dimensionality reduction operation on the input feature map so that the channel dimension c is half of the original; it is then processed by two stacked P-Bottlenecks, and the output features are residually concatenated with the original features to restore the dimension to c; then the SimAM attention mechanism module extracts key information from the concatenated features, and the last basic PConv module adjusts the features to a fixed output feature dimension to complete the feature output.
6. According to claim 5, a truck floor damage detection method based on improved YOLOv8 is characterized in that: The P-bottleneck module consists of two stacked basic PConv modules and residual connections. The convolution kernel size of each PConv module is set to 3, the stride is 1, and the residual connection used is an addition operation.
7. According to claim 5, a truck floor damage detection method based on improved YOLOv8 is characterized in that: The SimAM attention mechanism module is embedded after the P-Bottleneck module. The 3-D attention weight of the feature map is inferred in a parallel branch, a unique weight is assigned to each neuron, and the minimum energy function is derived to explore the importance of each neuron. The expression formula for enhancing the feature is: In the formula, X is the concatenated input feature map; sigmoid is the activation function to limit the maximum value of E; E is the value of all neurons in the attention feature map A collection of; represents the minimum energy, and its expression is: in, λ is a parameter and t is a neuron.