Floating object detection method based on improved YOLOv5s and transfer learning

Abstract

The invention discloses a floating object detection method based on improved YOLOv5s and transfer learning, and the method comprises the following steps: collecting floating object data, carrying out the preprocessing of the data, carrying out the enhancement of the data, and carrying out the normalization of the enhanced data; and then carrying out model improvement, OfficientNetv2 lightweight network structure replacement, CARAFE lightweight upsampling module change, replacement of PANet with a BiFPN structure, introduction of an attention mechanism, replacement of CIoU with an SIoU loss function, transfer learning, model evaluation and verification, and result analysis. Compared with a traditional target detection model, the method is improved in multiple aspects, the accuracy and precision of floating object detection can be improved by optimizing and improving the YOLOv5s model and combining transfer learning, the detection result is more reliable and accurate, the detection accuracy is higher compared with a traditional method, and the method is suitable for popularization and application. The YOLOv5s model is well balanced in the aspects of speed and precision, the performance can be further optimized by combining transfer learning, and the floating object detection efficiency and response speed are improved.

Description

Technical Field

[0001] The present invention relates to the technical field of river floating object detection, and in particular to a floating object detection method based on improved YOLOv5s and transfer learning. Background Art

[0002] In recent years, artificial intelligence has been developing continuously, and deep learning and machine learning technologies have been widely used in various fields. In this context, target detection algorithms have been widely used in various fields, among which the YOLO series of algorithms have attracted much attention due to their fast and efficient characteristics. However, in the field of water surface floating object detection, traditional target detection algorithms often face the challenges of interference from environmental factors such as illumination, occlusion, and noise, and accurate detection of floating objects is still a challenging problem.

[0003] In the existing technology, missed detection and false detection are prone to occur in the detection of small targets. The existing network models have problems such as low accuracy, weak model generalization ability, and large model size. Summary of the invention

[0004] In view of the above-mentioned problems existing in floating object detection, the present invention provides a floating object detection method based on improved YOLOv5s and transfer learning, which can achieve high-precision real-time detection of floating object targets on the water surface. During the test, no partial frame loss or detection frame offset occurs.

[0005] The solution adopted by the present invention to solve the technical problem is: a floating object detection method based on improved YOLOv5s and transfer learning, comprising the following steps:

[0006] S1. Collect floating object data;

[0007] S2, perform data preprocessing;

[0008] The collected floating object data were screened and the image size of floating objects on the water surface was uniformly cropped to 640×640;

[0009] S3. Data enhancement

[0010] The training set and test set are divided into a ratio of 9:1, and SAGAN is used to perform data enhancement on the image. In SAGAN, a non-local model is adopted to enable the generator and discriminator to effectively build the relationship between each region and directly calculate the relationship between two pixels.

[0011] S4. Data normalization

[0012] After using the LabelImg image annotation tool to select the YOLO format annotation, the TXT annotation file of the corresponding image is obtained;

[0013] S5, EfficientNetv2 lightweight network structure;

[0014] S6, CARAFE lightweight upsampling module

[0015] Change the original upsampling module of the YOLOv5s model to the CARAFE lightweight upsampling module. The CARAFE lightweight upsampling module consists of an upsampling kernel prediction module and a feature reorganization module. The feature map is input into the upsampling kernel prediction module.

[0016] S7, replace PANet with BiFPN structure;

[0017] S8, introduce the attention mechanism;

[0018] S9, replace CIoU with SIoU loss function;

[0019] S10. Transfer Learning

[0020] The transfer learning method is adopted, with the ImageNet dataset as the source domain and the river bank floating object dataset as the auxiliary domain. The model weight is transferred through pre-training as the initial weight parameter of the model to solve the problem of small sample data volume and improve the generalization ability of the model and the network training speed.

[0021] S11. Model evaluation and result analysis.

[0022] Furthermore, step S5 uses EfficientNetv2 to replace the shallow MBConv structure of EfficientNet with a Fused-MBConv structure, that is, the 1×1 dilated convolution and 3×3DW convolution of the main branch of the structure are replaced with a 3×3 ordinary convolution. EfficientNetv2 uses NAS to search for the best combination of MBConv and Fused-MBConv, and after balancing speed and performance, the EfficientNetv2-B0 structure is obtained.

[0023] Furthermore, in step S6, the number of channels C of the feature layer of the input feature map is compressed to C after a CBS convolution layer with a convolution kernel of 1×1. m , as shown below:

[0024] C m =σ 2 ·(K up ) 2

[0025] Where σ is the upsampling factor, K up is the upsampling kernel size;

[0026] For the input feature map after channel compression, use a convolution kernel size of ke ×k e The convolution layer is used to predict the upper convolution kernel, and the number of output channels is σ 2 K 2 up , and then use the PixelShuffle method to expand the channel dimension in the spatial dimension to obtain a shape of σH·σW·K 2 up The upsampling kernel is normalized using the Softmax activation function.

[0027] Furthermore, in step S7, a weighted feature fusion mechanism is introduced through BiFPN, as shown in the following formula:

[0028]

[0029] The value of ε is 0.0001, which makes the denominator not equal to 0 to ensure numerical stability. i The weight range is between [0,1], I i is the input feature map.

[0030] Furthermore, step S8 introduces the EMA attention mechanism, which is mainly composed of a 1×1 branch, a 3×3 branch and a cross-space learning module.

[0031] Further, the method of introducing the attention mechanism in step S8 includes:

[0032] S81. For any given input feature map, set

[0033] X∈R C×H×W

[0034] In the formula, C, H, and W are the number of channels, height, and width of the input feature map respectively;

[0035] EMA divides X into G sub-features along the cross-channel dimension for learning different semantics, where the group style can be defined as:

[0036] X=[X 0 ,X i ,…,X G-1 ],X i ∈R C∥G×H×W

[0037] In the formula, G<<C;

[0038] S82. In order to capture the dependencies between all channels and reduce computational overhead, EMA uses two one-dimensional global average pooling operations to encode channels along two spatial directions in the 1×1 branch, and only stacks one 3×3 kernel in the 3×3 branch to capture multi-scale feature representation. In EMA, a cross-spatial information aggregation method in different spatial dimensional directions can achieve richer feature aggregation. First, one of the two tensors is the output of the 1×1 branch and the other is the output of the 3×3 branch. Then, the output of the 1×1 branch is encoded with global spatial information using two-dimensional global average pooling. Before the channel feature joint activation mechanism, the output of the smallest branch is directly converted to the corresponding dimensional shape, that is:

[0039] R 1 1×C∥G ×R 3 C∥G×HW

[0040] The two-dimensional global pooling operation formula is:

[0041]

[0042] Where, X c Represents the input features of the Cth channel;

[0043] S83. In order to improve the computational efficiency of the model, Softmax is used at the output of the two-dimensional global average pooling to fit the above linear transformation. The first spatial attention map is obtained by multiplying the output of the above parallel processing with the matrix dot product operation. The 3×3 branch is also encoded with the global spatial information using the two-dimensional global average pooling. The 1×1 branch is directly converted to the corresponding dimensional shape before the channel feature joint activation mechanism, that is:

[0044] R 3 1×C∥G ×R 1 C∥G×HW

[0045] On this basis, a second spatial attention map that retains the entire precise spatial position information is derived. Finally, the output feature map within each group is calculated as a set of the two generated spatial attention weight values, and then the Sigmoid function is used to capture the pairwise relationship at the pixel level and highlight the global context of all pixels.

[0046] Further, in step S9, CIoU is replaced by SIoU loss function.

[0047] S91, SIoU loss function is composed of angle, distance, shape and intersection-over-union ratio (I IoU) The four cost functions are composed of the angle cost function, which is first predicted on the X-axis or Y-axis, and then approached along the X-axis or Y-axis. During the convergence process, if α≤π / 4, α is minimized, otherwise β is minimized;

[0048] The angle cost function Λ is as follows:

[0049]

[0050] The distance cost function Δ is as follows:

[0051]

[0052] In the formula, (b cx ,b cy ) is the center point of the prediction box, is the center point of the real frame, α is the sum of σ and C w Angle, C w , C h The difference between the horizontal and vertical coordinates of the center point of the predicted box and the center point of the real box, σ is the distance between the center point of the real box and the predicted box, w and h are the width and height of the predicted box;

[0053] S92. From the above formula, when α is close to 0, the distance cost contribution decreases; when α is close to π / 4, the distance cost function Δ is larger, and as the angle increases, γ is given time priority.

[0054] The shape cost function Ω is as follows:

[0055]

[0056] In the formula, w gt and h gt is the width and height of the real frame;

[0057] S93, The value defines the shape cost, and the SIoU loss function is defined The value range is [2,6], the box loss function L box for:

[0058]

[0059] The total loss function L is composed of the classification loss L cls and box loss L box The two components are as follows:

[0060] L=W box L box +W cls L cls

[0061] Where, L cls is the focal loss, W boxand W cls are the box and classification loss weights respectively.

[0062] Furthermore, step S11 is a method for model evaluation and result analysis: based on the improved YOLOv5 model of S5-S9, the transfer learning result of S10 is used as the initialization weight parameter, and then the model is trained through the training set, tested by the test set, and the model performance is evaluated; the accuracy, frame rate, F1 value and recall rate of the model before and after improvement are calculated to evaluate the prediction performance of the model; then, the accuracy, frame rate, F1 value and recall rate of the model before and after improvement are compared, and the Loss curve and accuracy curve comparison chart of the model before and after improvement are drawn to evaluate the effect of improved YOLOv5s and transfer learning on floating object target detection.

[0063] Beneficial effects of the present invention: In order to improve the accuracy and robustness of the detection of floating objects on the water surface, the present invention provides a floating object detection method that combines an improved YOLOv5s algorithm and transfer learning. The YOLOv5s algorithm has high detection accuracy and speed. By improving and optimizing it, it can better adapt to the detection task of floating objects on the water surface. At the same time, by using transfer learning, the model can be pre-trained on a large-scale data set, and then fine-tuned on the water surface floating object data set, to achieve targeted model optimization and improve the floating object detection effect.

[0064] Compared with the traditional target detection method, the floating object detection method of the present invention introduces more advanced algorithms and technical means in floating object detection, which can overcome environmental interference and improve accuracy and robustness. At the same time, combined with transfer learning, the information of large-scale data sets can be effectively utilized to improve model performance. This method of combining improved YOLOv5s and transfer learning has brought new technical directions and breakthroughs to the field of surface floating object detection.

[0065] The basic principle of the floating object detection method based on improved YOLOv5s and transfer learning of the present invention is to combine the YOLOv5s target detection algorithm and transfer learning technology. First, the YOLOv5s algorithm is improved and optimized, such as adjusting the network structure, adding feature fusion modules, introducing attention mechanisms, etc., to meet the needs of floating object detection tasks. Then, using transfer learning, the model is pre-trained on a large-scale data set, and then fine-tuned on the floating object data set to quickly and effectively improve the accuracy and generalization ability of the model in floating object detection. By combining the improved target detection model and transfer learning strategy, the purpose of more accurate and efficient detection of floating objects is achieved, providing cutting-edge technical support for surface floating object monitoring in practical applications.

[0066] The floating object detection method based on improved YOLOv5s and transfer learning described in the present invention can significantly improve the accuracy and precision of floating object detection by optimizing the improved model and the strategy of transfer learning, so that the detection results are more reliable and accurate, and have a higher detection accuracy rate than traditional methods. Secondly, the speed and accuracy balance of the YOLOv5s model is combined with the optimization of transfer learning, which can further improve the efficiency and response speed of floating object detection. In addition, transfer learning enables the model to learn more characteristic information of floating objects from existing data, improves the generalization ability of the model, and enables it to perform better in new environments. Most importantly, YOLOv5s has a relatively simple model architecture and training process, and the model after transfer learning is also easier to deploy and has strong practicality. In summary, the present invention provides an advanced technical support, which is expected to make significant progress and application effects in the field of floating object detection. BRIEF DESCRIPTION OF THE DRAWINGS

[0067] Figure 1 A box view of a floating object detection method based on improved YOLOv5s and transfer learning;

[0068] Figure 2 Loss curves before and after transfer learning for different models;

[0069] Figure 3 Accuracy curves before and after transfer learning for different models. DETAILED DESCRIPTION

[0070] In order to make the purpose, technical solution and advantages of the present invention clearer, the technical solution of the present invention will be described in detail below.

[0071] Example 1: The floating object detection method based on improved YOLOv5s and transfer learning provided by the present invention may include the following specific implementations:

[0072] Step 1, collect floating object data:

[0073] We selected rivers, lakes, ponds and other environments, and photographed the above four types of floating objects from different angles and in different weather conditions.

[0074] Step 2, data preprocessing:

[0075] The data was screened to remove images with extremely high similarity and images with insufficient clarity; the images of floating objects on the water surface were uniformly cropped to a size of 640×640.

[0076] Step 3, data enhancement:

[0077] The training set and test set are divided into a ratio of 9:1. SAGAN is used to enhance the image data. A non-local model is adopted in SAGAN, which enables the generator and discriminator to effectively build the relationship between each region and directly calculate the relationship between two pixels. The self-attention mechanism simply calculates the response of a single position in the weighted sum of the features of all positions. This mechanism allows the network to focus on areas that are scattered in different positions but have structural connections.

[0078] The principle is that in the generative network, the prior input noise and conditional information are combined to form a multi-mode vector, which is sent to the generative network as input representation information. The category label is introduced as a conditional variable in the discriminative network to judge not only whether the image is true or false, but also whether the generated category matches the category of the input image.

[0079] Step 4, data normalization:

[0080] After using the LabelImg image annotation tool to select the YOLO format annotation, the TXT annotation file of the corresponding image is obtained.

[0081] Step 5, improve the YOLOv5s model:

[0082] The lightweight network structure of EfficientNetv2 is replaced, the CARAFE lightweight upsampling module is changed, the PANet is replaced with the BiFPN structure, the attention mechanism is introduced, and the CIoU loss function is replaced with the SIoU loss function.

[0083] EfficientNetv2 is a lightweight CNN network improved by TAN and others based on EfficientNet. Compared with EfficientNet, EfficientNetv2 achieves faster training and reasoning speeds with similar accuracy and parameter count. EfficientNet uses a large number of depthwise (DW) convolutions. Although it has fewer parameters and computational complexity than ordinary convolutions, using DW convolutions in the shallow layers of the network will be very slow, making it difficult to fully utilize existing graphics processing units (GPUs).

[0084] To solve this problem, EfficientNetv2 replaces the shallow MBConv structure of EfficientNet with the Fused-MBConv structure, that is, replaces the 1×1 dilated convolution and 3×3DW convolution of the main branch of the structure with a 3×3 ordinary convolution. EfficientNetv2 uses NAS technology to search for the best combination of MBConv and Fused-MBConv, and obtains the EfficientNetv2-B0 structure after balancing speed and performance. The numbers after MBConv and Fused-MBConv indicate the feature channel expansion multiples. The shallow Fused-MBConv structure does not use the SE module in the actual construction process of the EfficientNetv2-B0 structure network.

[0085] Step 6, transfer learning:

[0086] The transfer learning method is adopted, with the ImageNet dataset as the source domain and the river bank floating object dataset as the auxiliary domain. The model weight transfer is realized through pre-training as the initialization weight parameter of the model to solve the problem of small sample data volume and improve the generalization ability of the model and the network training speed.

[0087] The upsampling operation can be interpreted as the upsampling kernel at each position performing feature reorganization by dot product with the pixels in the corresponding neighborhood in the input feature map. The most commonly used upsampling methods are the nearest neighbor and bilinear interpolation. These two methods only determine the upsampling kernel through the local neighborhood of the pixel point, do not use the global semantic information of the feature map, and have a small receptive field, which cannot effectively extract the global features of the image.

[0088] In order to better utilize the global features around the pixels during the upsampling process, the original upsampling module of the YOLOv5s model is changed to a CARAFE lightweight upsampling module in the present invention. The CARAFE lightweight upsampling module consists of an upsampling kernel prediction module and a feature reorganization module. The feature map is first passed to the upsampling kernel prediction module, and the number of channels C in the feature layer is reduced to C after a CBS convolution layer with a convolution kernel of 1×1. m , as shown in formula (1):

[0089] C m =σ 2 ·(K up ) 2 (1)

[0090] Where σ is the upsampling factor, K up is the upsampling kernel size.

[0091] For the input feature map after channel compression, use a convolution kernel size of k e ×k eThe convolution layer is used to predict the upper convolution kernel, and the number of output channels is σ 2 K 2 up , and then use the PixelShuffle method to expand the channel dimension in the spatial dimension to obtain a shape of σH·σW·K 2 up The upsampling kernel is normalized using the Softmax activation function.

[0092] Step 7, replace PANet with BiFPN structure:

[0093] In the floating object dataset, some floating objects are small in size. Although YOLOv5s has a PANet structure for feature extraction, there are still omissions and misjudgments. In order to improve the detection accuracy of the model for small floating objects, BiFPN is introduced to replace the original PANet structure.

[0094] The step S7 specifically includes the following process:

[0095] S71, in the PANet structure, there is only one top-down path and one bottom-up path. Unlike PANet, BiFPN uses each bidirectional path as a feature network layer and stacks it multiple times to achieve higher-level feature fusion. At the same time, BiFPN adds an additional path between the original input node and the output node to fuse more features without increasing too much computational cost.

[0096] S72, the feature fusion in PANet is just a simple addition of feature maps, and the feature maps are not differentiated. However, different input feature maps have different resolutions and contribute differently to the fused input feature maps, so simple addition cannot achieve the best effect. For this reason, BiFPN introduces a weighted feature fusion mechanism, as shown in formula (2):

[0097]

[0098] The value of ε is 0.0001, which makes the denominator not equal to 0 to ensure numerical stability. i The weight range is between [0,1], I i is the input feature map.

[0099] Step 8, introduce the attention mechanism:

[0100] An efficient multi-scale attention mechanism (EMA) is introduced into the Backbone part of YOLOv5s to increase the weight of important features and reduce the weight of irrelevant features such as the matrix and the edge of the hole tray, thereby improving the detection accuracy of the model.

[0101] Specifically, the S81 and EMA attention mechanisms are mainly composed of a 1×1 branch, a 3×3 branch, and a cross-space learning module. For any given input feature map, we assume that:

[0102] X∈R C×H×W (3)

[0103] Where C, H, and W are the number of channels, height, and width of the input feature map, respectively.

[0104] EMA divides X into G sub-features along the cross-channel dimension for learning different semantics, where the group style can be defined as:

[0105] X=[X 0 ,X i ,…,X G-1 ],X i ∈R C∥G×H×W (4)

[0106] In the formula, G<<C.

[0107] S82. In order to capture the dependencies between all channels and reduce computational overhead, EMA uses two one-dimensional global average pooling operations to encode the channels along two spatial directions in the 1×1 branch, and only stacks one 3×3 kernel in the 3×3 branch to capture multi-scale feature representation. In EMA, a cross-spatial information aggregation method in different spatial dimensional directions can achieve richer feature aggregation. First, one of the two tensors is the output of the 1×1 branch and the other is the output of the 3×3 branch. Then, the output of the 1×1 branch is encoded with global spatial information using two-dimensional global average pooling, and the output of the smallest branch is directly converted to the corresponding dimensional shape before the channel feature joint activation mechanism, that is:

[0108] R 1 1×C∥G ×R 3 C∥G×HW (5)

[0109] The two-dimensional global pooling operation formula is:

[0110]

[0111] Where, X c Represents the input features of the Cth channel.

[0112] S83. In order to improve the computational efficiency of the model, Softmax is used at the output of the two-dimensional global average pooling to fit the above linear transformation. The first spatial attention map is obtained by multiplying the output of the above parallel processing with the matrix dot product operation. The 3×3 branch is also encoded with the global spatial information using the two-dimensional global average pooling, and the 1×1 branch is directly converted to the corresponding dimensional shape before the channel feature joint activation mechanism, that is:

[0113] R 3 1×C∥G ×R 1 C∥G×HW (7)

[0114] On this basis, a second spatial attention map is derived that retains the entire precise spatial position information. Finally, the output feature map within each group is calculated as a set of the two generated spatial attention weight values, and then the Sigmoid function is used to capture the pairwise relationship at the pixel level and highlight the global context of all pixels.

[0115] Step 9, replace CIoU with SIoU loss function:

[0116] Although the CIoU loss function in YOLOv5s takes into account the aspect ratio of the regression box as the loss function of the regression box scale, it ignores the problem of direction mismatch between the predicted box and the true box, which will lead to slower model convergence and lower efficiency. In order to optimize the CIoU loss function, the present invention adopts the SIoU loss function.

[0117] The specific steps include:

[0118] S91, SIoU loss function is composed of angle, distance, shape and intersection-over-union ratio (I IoU )These four cost functions are constructed. The angular cost function is first predicted on the X-axis or Y-axis, and then continues to approach along that axis. During the convergence process, if α≤π / 4, α is minimized, otherwise β is minimized.

[0119] The angle cost function Λ is as follows:

[0120]

[0121] The distance cost function Δ is as follows:

[0122]

[0123] In the formula, (b cx ,b cy ) is the center point of the prediction box, is the center point of the real frame, α is the sum of σ and C w Angle, C w , C hThe difference between the horizontal and vertical coordinates of the center point of the predicted box and the center point of the real box, σ is the distance between the center point of the real box and the predicted box, and w and h are the width and height of the predicted box.

[0124] S92, from the above formula, when α is close to 0, the distance cost contribution decreases. When α is close to π / 4, the distance cost function Δ is larger. As the angle increases, γ is given time priority.

[0125] The shape cost function Ω is as follows:

[0126]

[0127] In the formula, w gt and h gt is the width and height of the real frame

[0128] S93, Value defines shape cost, experimental The value is close to 4, SIoU loss function definition The value range is [2,6]. Box loss function L box for:

[0129]

[0130] The total loss function L is composed of the classification loss L cls and box loss L box The two components are as follows:

[0131] L=W box L box +W cls L cls (15)

[0132] Where, L cls is the focal loss, W box and W cls are the box and classification loss weights respectively.

[0133] Step 10, transfer learning:

[0134] The transfer learning method is adopted, with the ImageNet dataset as the source domain and the river bank floating object dataset as the auxiliary domain. The model weight transfer is realized through pre-training as the initialization weight parameter of the model to solve the problem of small sample data volume and improve the generalization ability of the model and the network training speed.

[0135] Step 11, model evaluation and result analysis:

[0136] Based on the improved YOLOv5 model of S5-S9, the transfer learning results of S10 are used as the initial weight parameters, and then the model is trained through the training set, tested with the test set, and the model performance is evaluated. The accuracy, frame rate, F1 value and recall rate of the model before and after the improvement are calculated to evaluate the prediction performance of the model. Then, the accuracy, frame rate, F1 value and recall rate of the model before and after the improvement are compared, and the loss curve and accuracy curve comparison chart of the model before and after the improvement are drawn to evaluate the effect of improved YOLOv5s and transfer learning on floating object target detection.

[0137] The accuracy of the improved YOLOv5 model of the present invention is improved by 4.6 percentage points, the F1 index is improved by 3.9 percentage points, and the recall rate is improved by 0.6 percentage points. At the same time, the frame rate can still reach a detection speed of 37.2 frames per second, which is higher than the 35.6 frames per second before the improvement. Then, the Loss curve and accuracy curve comparison of the model before and after the improvement are drawn, as shown in the figure below: Figure 2 , 3 As shown in the figure, during the entire training phase of the model, the improved model has a lower loss value, higher accuracy, and faster model convergence. Overall, the improved YOLOv5s model performs better in floating object detection after transfer learning than the unimproved YOLOv5s model.

[0138] The above specific implementation methods can enable the improved YOLOv5s model to be implemented in the present invention, providing a more sufficient data basis for the training of the floating object target detection model.

[0139] The invention discloses a method for detecting floating objects based on improved YOLOv5s and transfer learning. Collect floating object data, pre-process the data, enhance the data, and normalize the enhanced data; then improve the model, replace the EfficientNetv2 lightweight network structure, change the CARAFE lightweight upsampling module, replace the PANet with the BiFPN structure, introduce the attention mechanism, replace the CIoU with the SIoU loss function, transfer learning, model evaluation and verification, and result analysis. This invention brings multiple benefits: first, by optimizing the improved model and the strategy of transfer learning, the accuracy and precision of floating object detection can be significantly improved, making the detection results more reliable and accurate, and having a higher detection accuracy rate than traditional methods. Secondly, combining the speed and accuracy balance of the YOLOv5s model, coupled with the optimization of transfer learning, the efficiency and response speed of floating object detection can be further improved. In addition, transfer learning enables the model to learn more characteristic information of floating objects from existing data, improves the generalization ability of the model, and enables it to perform better in new environments. Most importantly, YOLOv5s has a relatively simple model architecture and training process. The model after transfer learning is also easier to deploy and has strong practicality. These advantages work together to provide an advanced technical support. In the field of floating object detection, this method has broad application prospects, helps to improve the training effect of floating object detection models, reduces the cost of data collection and labeling, and provides a more comprehensive and representative data foundation to promote the development of the field of floating object detection.

[0140] Obviously, the described embodiments are only some embodiments of the present invention, not all embodiments. Based on the embodiments of the present invention, all other implementation methods obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

Claims

1. A floating object detection method based on improved YOLOv5s and transfer learning, characterized in that: The following steps are involved: S1. Collect floating object data; S2, perform data preprocessing; The collected floating object data were screened and the image size of floating objects on the water surface was uniformly cropped to 640×640. S3. Data enhancement The training set and test set are divided into a ratio of 9:1, and SAGAN is used to perform data augmentation on the images. In SAGAN, a non-local model is adopted to enable the generator and discriminator to effectively build the relationship between each region and directly calculate the relationship between two pixels. S4. Data normalization After using the Labelmg image annotation tool to select the YOLO format annotation, the TXT annotation file of the corresponding image is obtained; S5, EfficientNetv2 lightweight network structure; S6, CARAFE lightweight upsampling module Change the original upsampling module of the YOLOv5s model to the CARAFE lightweight upsampling module. The CARAFE lightweight upsampling module consists of an upsampling kernel prediction module and a feature reorganization module. The feature map is input into the upsampling kernel prediction module. S7, replace PANet with BiFPN structure; S8, introduce the attention mechanism; S9, replace CIoU with SIoU loss function; S10. Transfer Learning The transfer learning method is adopted, with the ImageNet dataset as the source domain and the river bank floating object dataset as the auxiliary domain. The model weight is transferred through pre-training as the initial weight parameter of the model to solve the problem of small sample data volume and improve the generalization ability of the model and the network training speed. S11. Model evaluation and result analysis.

2. The floating object detection method based on improved YOLOv5s and transfer learning according to claim 1, characterized in that: Step S5 uses EfficientNetv2 to replace the shallow MBConv structure of EfficientNet with the Fused-MBConv structure, that is, replace the 1×1 dilated convolution and 3×3 DW convolution of the main branch of the structure with a 3×3 ordinary convolution. EfficientNetv2 uses NAS to search for the best combination of MBConv and Fused-MBConv, and after balancing speed and performance, obtains the EfficientNetv2-B0 structure.

3. The floating object detection method based on improved YOLOv5s and transfer learning according to claim 1, characterized in that: In step S6, the number of channels C of the feature layer of the input feature map is compressed to Cm after a CBS convolution layer with a convolution kernel of 1×1, as shown in the following formula: C m =s 2 ·(K up ) 2 Where σ is the upsampling factor, K up is the upsampling kernel size; For the input feature map after channel compression, use a convolution kernel size of k e ×k e The convolution layer is used to predict the upper convolution kernel, and the number of output channels is σ 2 K 2 up , and then use the PixelShuffle method to expand the channel dimension in the spatial dimension to obtain a shape of σH·σW·K 2 up The upsampling kernel is normalized using the Softmax activation function.

4. The floating object detection method based on improved YOLOv5s and transfer learning according to claim 1, characterized in that: In step S7, a weighted feature fusion mechanism is introduced through BiFPN, as shown in the following formula: The value of ε is 0.0001, which makes the denominator not equal to 0 to ensure numerical stability. i The weight range is between [0, 1], I i is the input feature map.

5. The floating object detection method based on improved YOLOv5s and transfer learning according to claim 1, characterized in that: Step S8 introduces the EMA attention mechanism, which is mainly composed of a 1×1 branch, a 3×3 branch and a cross-space learning module.

6. The floating object detection method based on improved YOLOv5s and transfer learning according to claim 5, characterized in that: The method of introducing the attention mechanism in step S8 includes: S81, for any given input feature map, setting X∈R C×H×W Where C, H, and W are the number of channels, height, and width of the input feature map, respectively; EMA divides X into G sub-features along the cross-channel dimension for learning different semantics, where the group style can be defined as: X=[X0,X i ,…,X G-1 ],X i ∈R C / / G×H×W In the formula, G<<C; S82. In order to capture the dependencies between all channels and reduce computational overhead, EMA uses two one-dimensional global average pooling operations to encode channels along two spatial directions in the 1×1 branch, and only stacks one 3×3 kernel in the 3×3 branch to capture multi-scale feature representation. In EMA, a cross-spatial information aggregation method in different spatial dimensional directions can achieve richer feature aggregation. First, one of the two tensors is the output of the 1×1 branch and the other is the output of the 3×3 branch. Then, the output of the 1×1 branch is encoded with global spatial information using two-dimensional global average pooling. Before the channel feature joint activation mechanism, the output of the smallest branch is directly converted to the corresponding dimensional shape, that is: R1 1×C / / G ×R3 C / / G×HW The two-dimensional global pooling operation formula is: Where, X c Represents the input features of the Cth channel; S83. In order to improve the computational efficiency of the model, Softmax is used at the output of the two-dimensional global average pooling to fit the above linear transformation. The first spatial attention map is obtained by multiplying the output of the above parallel processing with the matrix dot product operation. The 3×3 branch is also encoded with the global spatial information using the two-dimensional global average pooling. The 1×1 branch is directly converted to the corresponding dimensional shape before the channel feature joint activation mechanism, that is: R3 1×C / / G ×R1 C / / G×HW On this basis, a second spatial attention map that retains the entire precise spatial position information is derived. Finally, the output feature map within each group is calculated as a set of the two generated spatial attention weight values, and then the Sigmoid function is used to capture the pairwise relationship at the pixel level and highlight the global context of all pixels.

7. The floating object detection method based on improved YOLOv5s and transfer learning according to claim 1, characterized in that: Step S9 replaces CIoU with SIoU loss function. S91, SIoU loss function is composed of angle, distance, shape and intersection-over-union ratio (I IoU ) The four cost functions are composed of the angle cost function, the angle cost function is first predicted on the X-axis or Y-axis, and then approached along the X-axis or Y-axis. During the convergence process, if α≤π / 4, α is minimized, otherwise β is minimized; the angle cost function Λ is as follows: The distance cost function Δ is as follows: In the formula, (b cx , b cy ) is the center point of the prediction box, is the center point of the real frame, α is the sum of σ and C w Angle, C w , C h The difference between the horizontal coordinate and the vertical coordinate of the center point of the predicted box and the center point of the real box, σ is the distance between the center point of the real box and the predicted box, w and h are the width and height of the predicted box; S92. From the above formula, when α is close to 0, the contribution of the distance cost decreases; when α is close to π / 4, the larger the distance cost function Δ is, as the angle increases, γ is given time priority, The shape cost function Ω is as follows: Where W gt and h gt is the width and height of the real frame; S93, θ value defines the shape cost, SIoU loss function defines θ value range as [2, 6], box loss function L box for: The total loss function L is composed of the classification loss L cls and box loss L box The two components are as follows: L=W box L box +W cls L cls Where, L cls is the focal loss, W box and W cls are the box and classification loss weights respectively.

8. The floating object detection method based on improved YOLOv5s and transfer learning according to claim 1, characterized in that: Step S11: Model evaluation and result analysis method: Based on the improved YOLOv5 model of S5-S9, the transfer learning result of S10 is used as the initial weight parameter, and then the model is trained through the training set, tested with the test set, and the model performance is evaluated; the accuracy, frame rate, F1 value and recall rate of the model before and after the improvement are calculated to evaluate the prediction performance of the model; Then, the accuracy, frame rate, F1 value and recall rate of the model before and after improvement are compared, and the loss curve and accuracy curve comparison chart of the model before and after improvement are drawn to evaluate the effect of improved YOLOv5s and transfer learning on floating object detection.

Images