An underwater target detection method based on YOLOv5

Abstract

The present invention discloses a method for underwater target detection based on YOLOv5, comprising the following steps: obtaining an initial dataset of specific underwater targets; enhancing the initial dataset to increase the number of the initial dataset and obtain an expanded dataset; improving some C3 modules in the backbone network using the CA module; adding the SE attention module to a specific position in the backbone network; detecting the input image using a detection head based on anchor-free boxes; inputting the expanded dataset into a YOLOv5 model for training; and detecting underwater images using the trained YOLOv5 model. The present invention can improve the detection accuracy of the detection model for multiple specific targets in complex underwater environments without significantly reducing the detection speed.

Description

Technical Field

[0001] The present invention relates to the field of target detection technology, and in particular to an underwater target detection method based on YOLOv5. Background Art

[0002] Most existing target detection methods are primarily applied to terrestrial environments. Due to the convenience of land, the quality and quantity of images collected on land are ideal, allowing the training of target detection models to achieve the desired results. However, unlike terrestrial environments, underwater environments are often very complex (such as turbid water, uneven lighting, and interference from ocean currents). Multiple factors affect the quality and quantity of underwater images, which in turn affects the accuracy of underwater target detection. Considering the complexity of underwater environments, conventional detection models often need to be appropriately modified to enable them to better perform various target detection tasks in underwater environments. Summary of the Invention

[0003] The technical problem to be solved by the present invention is to provide an underwater target detection method based on YOLOv5, which can ensure that the detection accuracy of the YOLOv5 model for multiple specific targets in a complex underwater environment is improved without significantly reducing the detection speed.

[0004] To solve the above technical problems, the technical solution adopted by the present invention is: an underwater target detection method based on YOLOv5, comprising the following steps:

[0005] Step S1, obtaining an initial data set of specific underwater targets: the initial data set is obtained by marking underwater images using marking software, and marking various targets in the image with specific boxes;

[0006] Step S2: enhancing the initial data set to increase the amount of the initial data set to obtain an expanded data set;

[0007] Step S3: Introducing the CA module into the C3 module in the backbone feature extraction network of the YOLOv5 model to improve the detection performance of the YOLOv5 model for underwater blurred targets, and adding the SE attention module at the output end of the C3 module in the backbone feature extraction network to improve the detection accuracy of the YOLOv5 model for blurred objects in the image;

[0008] Step S4: input the expanded data set into the YOLOv5 model backbone feature extraction network to obtain a feature map. The output feature map of the SE attention module is input into the neck network for feature fusion. The feature map is bidirectionally fused with low-level spatial features and high-level semantic features using the feature pyramid network and the path aggregation network. The fused feature map is input into the detection network of the YOLOv5 model.

[0009] Step S5: The anchor-free detection box in the YOLOv5 model detection network processes the input fused feature map and identifies the type, location, and confidence of the detected target in the input image.

[0010] Step S6: input the expanded data set into the YOLOv5 model and set the corresponding number of training times. As the number of training times increases, the loss function curve of the YOLOv5 model gradually converges. When the loss function curve converges and stabilizes, the YOLOv5 model is trained to the optimal state and a corresponding weight file is generated.

[0011] Step S7: Download the trained weight file to the YOLOv5 model to detect the actual underwater image, mark the detected type, location, confidence level and other information in the image, and output the relevant file.

[0012] A further improvement of the technical solution of the present invention is that: in step S1, the acquired underwater image is labeled using the LabelImg program, its format is modified to the YOLO format, and applied to the training of the YOLOv5 model.

[0013] A further improvement of the technical solution of the present invention is that: in step S2, the initial data set is enhanced by using a mirror flipping method. The specific process is: setting the processing methods of left and right, up and down direction mirroring, left and right direction mirroring, and up and down direction mirroring respectively. Each time the flipping method is set, the initial data set is processed once, and finally an expanded data set is obtained. The expanded data set is three times the size of the initial data set.

[0014] A further improvement of the technical solution of the present invention is that the C3 module in the backbone feature extraction network of the YOLOv5 model in step S3 introduces the CA module by adding a CA module between the convolution module and the bottleneck module in the backbone branch of the C3-True module to improve the model's ability to extract shallow features.

[0015] A further improvement of the technical solution of the present invention is that: in the step S3, the C3 module in the backbone feature extraction network of the YOLOv5 model introduces the CA module by using the CA module to replace the bottleneck module in the backbone branch of the C3-False module.

[0016] A further improvement of the technical solution of the present invention is that in step S4, three types of feature maps are output in the neck network, corresponding to the detection of large, medium and small targets in the image, respectively, and the three features are input into the detection network for actual detection.

[0017] A further improvement of the technical solution of the present invention is that the detection network of the YOLOv5 model in step S5 consists of three detection layers, and its output is a feature map with dimensions of 80×80, 40×40 and 20×20, respectively, which is used to detect large, medium and small targets in the image.

[0018] A further improvement of the technical solution of the present invention is that the specific process of step S6 is: the expanded data set includes a training set, a validation set and a test set, and the ratio of the training set, the validation set and the test set is 8:1:1.

[0019] A further improvement of the technical solution of the present invention is that the loss function in step S6 is a CIOU loss function, and the specific calculation formula is:

[0020]

[0021] L CIOU =1-CIOU,

[0022] where ρ 2 (b,b gt ) represents the Euclidean distance between the center coordinates of the predicted box and the center coordinates of the real box, c represents the diagonal length of the rectangle circumscribing the two boxes, and αv represents the aspect ratio factor of the real box and the predicted box.

[0023] Due to the adoption of the above technical solution, the technical advancements achieved by the present invention are:

[0024] 1. This paper proposes an underwater target detection method based on YOLOv5 to detect targets in different scenarios. This method effectively solves the problem of reduced target detection accuracy due to the complex underwater environment. At the same time, it can ensure that the detection accuracy of the YOLOv5 model for various specific targets in complex underwater environments is improved without significantly reducing the detection speed.

[0025] 2. The present invention combines the CA module and the C3 module to form two new improved modules to replace some C3 modules in the original model, thereby improving the detection accuracy of the model while reducing the number of model parameters; the SE module is added to the output end of some C3 modules in the backbone network, so that the YOLOv5 model can better detect targets in blurred images; the detection box based on anchor boxes in the original detection network is changed to one based on no anchor boxes, which greatly reduces the time consumption and computing power required. BRIEF DESCRIPTION OF THE DRAWINGS

[0026] Figure 1 This is a structural diagram of the C3-1 module in the backbone network of the present invention;

[0027] Figure 2This is a structural diagram of the C3-2 module in the backbone network of the present invention. DETAILED DESCRIPTION

[0028] The present invention is described in further detail below in conjunction with the embodiments:

[0029] The underwater target detection method of the embodiment includes:

[0030] Step S1: Acquire an initial data set of specific underwater targets; the initial data set is obtained by marking underwater images using marking software, and marking various targets in the image with specific boxes.

[0031] The initial dataset of underwater specific targets comes from the underwater target detection group of the 2019 China Underwater Robot Competition. This initial dataset contains a total of 3701 underwater images. In this example, the LabelImg program is used to annotate the acquired underwater images, and the format is modified to YOLO format for application in YOLOv5 model training.

[0032] Step S2: Use data enhancement methods to enhance the initial dataset to increase the number of initial datasets and obtain an expanded dataset. The specific methods include:

[0033] The mirror flip method is used to set the left and right up and down direction mirroring, left and right direction mirroring, and up and down direction mirroring processing methods respectively. Each time the flip method is set, the initial dataset is processed once, and finally an expanded dataset is obtained. The number of the expanded dataset is 3 times that of the initial dataset.

[0034] Step S3: Before inputting the expanded dataset into the YOLOv5 model backbone network, appropriate measures are taken to modify the original backbone network of the YOLOv5 model to improve the backbone network's ability to extract features, thereby obtaining a richer feature map.

[0035] Specifically include:

[0036] 1) The CA module is combined with the C3 module at a specific position in the YOLOv5 model to form a new improved module to replace the original C3 module in the model, so as to improve the detection performance of the detection model for underwater blurred targets.

[0037] The first combination method is to add a CA module between the convolution module and the bottleneck module in the main branch of the C3-True module to improve the model's ability to extract shallow features. The second combination method is to use the CA module to replace the bottleneck module in the main branch of the C3-False module. The schematic diagram of the first combination method is as follows Figure 1As shown, it is named C3-1 module. The schematic diagram of the second combination mode is as follows Figure 2 As shown, it is named C3-2 module.

[0038] 2) Due to the complex underwater environment, the collected data sets are often fuzzy, so the present invention embeds the SE attention module in the backbone network of the YOLOv5 model.

[0039] This paper adds an SE attention module to the output of the C3 module in the backbone network to improve the detection model's accuracy for blurred objects in images. Embedding the SE attention module in the backbone network also improves the fusion feature layer in the neck network. While the original YOLOv5 model integrates the output feature maps of the C3 module in the backbone network into the neck network for feature fusion, this paper also integrates the output feature maps of the SE attention module into the neck network for feature fusion.

[0040] Step S4: The expanded dataset is input into the modified YOLOv5 model backbone feature extraction network to obtain a feature map. The output feature map of the SE attention module is input into the neck network for feature fusion. The feature pyramid network and the path aggregation network are used to bidirectionally fuse the low-level spatial features with the high-level semantic features of the feature map, and the fused feature map is input into the detection network.

[0041] The neck network is a series of feature aggregation layers that combine image features, primarily used to generate feature pyramids and path clustering networks. By fusing different features through bidirectional fusion, it strengthens the bottom-up path and improves the transmission of low-level features. The neck network outputs three types of feature maps, corresponding to the detection of large, medium, and small objects in the image. These three types of features are then fed into the detection network for actual detection.

[0042] Step S5: Modify the anchor box-based detection frame in the YOLOv5 model detection network to a detection frame based on no anchor box, use the detection frame to process the input fused features, and identify the type, location and confidence of the detected target in the input image.

[0043] Since the setting of anchor boxes needs to be manually designed, different anchor boxes need to be designed for different data sets, and the number of anchor boxes is huge, a lot of computing power is required when calculating IOU, and the detection box is modified to a module based on no anchor boxes.

[0044] The detection network of the YOLOv5 model consists of three detection layers, whose outputs are feature maps with dimensions of 80×80, 40×40, and 20×20, respectively, which are used to detect large, medium, and small objects in the image.

[0045] Step S6: Input the expanded data set into the YOLOv5 model and set the corresponding number of training times. As the number of training times increases, the loss function curve of the YOLOv5 model gradually converges. When the loss function curve converges and stabilizes, the YOLOv5 model is trained to the optimal state and generates the corresponding weight file.

[0046] Divide the prepared expanded dataset into three parts according to the ratio of training set: validation set: test set equal to 8:1:1. Build the corresponding training environment in a Linux environment. The specific training environment is: Python3.8, CUDA11.1, pytorch1.8.0, and upload the expanded dataset and the modified YOLOv5 model to the environment. Set the corresponding number of training times. As the number of training times increases, the loss function curve of the YOLOv5 model gradually converges. When the loss function curve converges and stabilizes, the YOLOv5 model is trained to the optimal level. The loss function is the CIOU loss function, and the specific calculation formula is:

[0047]

[0048] L CIOU =1-CIOU,

[0049] where ρ 2 (b,b gt ) represents the Euclidean distance between the center coordinates of the predicted box and the center coordinates of the real box, c represents the diagonal length of the rectangle circumscribing the two boxes, and αv represents the aspect ratio factor of the real box and the predicted box.

[0050] Step S7: Download the trained weight file to the YOLOv5 model, input a new underwater picture, and the YOLOv5 model detects it, annotates the detected type, location, confidence level and other information in the picture, and outputs the relevant file.

[0051] The YOLOv5-based underwater target detection method of this embodiment can effectively detect targets in various scenarios. This method effectively solves the problem of reduced target detection accuracy due to complex underwater environments. Furthermore, this embodiment can improve the detection accuracy of the YOLOv5 model for various specific targets in complex underwater environments without significantly reducing detection speed.

Claims

1. A method for underwater target detection based on YOLOv5, characterized by: The steps include: Step S1, obtaining an initial data set of underwater targets: the initial data set is obtained by marking underwater images using marking software, and marking various targets in the image with boxes; Step S2: enhancing the initial data set to increase the amount of the initial data set to obtain an expanded data set; Step S3, the C3 module in the trunk feature extraction network of the YOLOv5 model introduces a CA module to improve the detection performance of the YOLOv5 model for underwater blurred targets, and adds an SE attention module at the output end of the C3 module in the trunk feature extraction network to improve the detection accuracy of the YOLOv5 model for blurred objects in the image; wherein, the C3 module in the trunk feature extraction network of the YOLOv5 model introduces the CA module by adding a CA module between the convolution module and the bottleneck module in the trunk branch of the C3-True module to improve the model's ability to extract shallow features; Step S4: input the expanded data set into the YOLOv5 model backbone feature extraction network to obtain a feature map. The output feature map of the SE attention module is input into the neck network for feature fusion. The feature map is bidirectionally fused with low-level spatial features and high-level semantic features using the feature pyramid network and the path aggregation network. The fused feature map is input into the detection network of the YOLOv5 model. Step S5: The anchor-free detection box in the YOLOv5 model detection network processes the input fused feature map and identifies the type, location, and confidence of the detected target in the input image. Step S6: input the expanded data set into the YOLOv5 model and set the corresponding number of training times. As the number of training times increases, the loss function curve of the YOLOv5 model gradually converges. When the loss function curve converges and stabilizes, the YOLOv5 model is trained to the optimal state and a corresponding weight file is generated. Step S7: Download the trained weight file to the YOLOv5 model to detect the actual underwater image, mark the detected type, location, and confidence information in the image, and output the relevant file.

2. The underwater target detection method based on YOLOv5 according to claim 1, characterized in that: In step S1, the LabelImg program is used to label the acquired underwater image, and its format is modified to the YOLO format, which is applied to the training of the YOLOv5 model.

3. The underwater target detection method based on YOLOv5 according to claim 1, characterized in that: In step S2, the initial data set is enhanced by using a mirror flipping method. The specific process is: setting the processing methods of left and right, up and down direction mirroring, left and right direction mirroring, and up and down direction mirroring respectively. Each time the flipping method is set, the initial data set is processed once, and finally an expanded data set is obtained. The expanded data set is three times the size of the initial data set.

4. The underwater target detection method based on YOLOv5 according to claim 1, characterized in that: In step S3, the C3 module in the backbone feature extraction network of the YOLOv5 model introduces the CA module by replacing the bottleneck module in the backbone branch of the C3-False module with the CA module.

5. The underwater target detection method based on YOLOv5 according to claim 1, characterized in that: In step S4, three types of feature maps are output in the neck network, corresponding to the detection of large, medium and small targets in the image, respectively, and the three features are input into the detection network for actual detection.

6. The underwater target detection method based on YOLOv5 according to claim 5, characterized in that: The detection network of the YOLOv5 model in step S5 consists of three detection layers, and its output is a feature map with dimensions of 80×80, 40×40, and 20×20, respectively, which is used to detect large, medium, and small objects in the image.

7. The underwater target detection method based on YOLOv5 according to claim 1, characterized in that: The expanded data set in step S6 includes a training set, a validation set, and a test set, and the ratio of the training set, the validation set, and the test set is 8:1:

1.

8. The underwater target detection method based on YOLOv5 according to claim 1, characterized in that: The loss function in step S6 is the CIOU loss function, and the specific calculation formula is: L CIOU =1-CIOU, where ρ 2 (b,b gt ) represents the Euclidean distance between the center coordinates of the predicted box and the center coordinates of the real box, c represents the diagonal length of the rectangle circumscribing the two boxes, and αv represents the aspect ratio factor of the real box and the predicted box.

Images