Underwater target detection method based on mask self-supervised contrast learning

Abstract

The invention discloses an underwater target detection method based on mask self-supervised contrast learning, relates to the field of underwater target detection, and aims to solve the problems of low recognition precision and poor scene adaptability in the existing underwater target detection method. The detection method specifically comprises the following steps: constructing an underwater target detection reference data set; the Swinin-Transform and the adaptive feature fusion pyramid module are added to a comparative learning framework, and a mask is added to the image for pre-training; freezing backbone network parameters, and adjusting adaptive feature fusion pyramid module parameters; a backbone network in the YOLOX is replaced with a Swinin-Transform, and an improved adaptive feature fusion pyramid module is used; and performing underwater target detection by using the improved YOLOX to obtain the position and category information of the underwater target. The target detection method for the underwater scene is improved, the structure of the network is improved, the parameters of the network are optimized, the adaptive capacity of the target detection algorithm in the underwater scene is improved, and the method is suitable for high-precision target detection in the underwater complex environment.

Description

Technical Field

[0001] The present invention relates to the fields of computer technology, deep learning, and underwater machine vision, and in particular to the field of underwater target detection, and specifically to an underwater target detection method based on mask self-supervised contrast learning. Background Art

[0002] The ocean is the birthplace of life and holds a vast reservoir of exploitable resources. Marine exploration is crucial for the rational utilization of resources and the sustainable development of the ecological environment. In recent years, with the advancement of deep learning technology, underwater unmanned vehicles (UAVs) equipped with intelligent object detection algorithms have played a significant role in marine resource exploration. However, the complex underwater environment results in high noise levels in captured images, making it difficult for object detection networks to learn good visual representations. This poses a significant challenge to common object detection algorithms. Therefore, improving the quality of the visual representations learned by object detection networks and designing high-precision underwater object detection algorithms have become key research areas.

[0003] Self-supervised learning is a type of unsupervised learning, mainly including the Masked Image Modeling (MIM) method and the Instance Discriminant (ID) method. Its main purpose is to enable the network to learn richer semantic representations. It can better complete related tasks by mining the structure or features contained in the data itself. However, the ID method avoids features but lacks spatial sensitivity, while the MIM method maintains spatial sensitivity but does not perform semantic alignment. Therefore, combining the advantages of the MIM method and the ID method, self-supervised learning is introduced into underwater target detection to adapt to the complex and changing underwater environment and learn more representative visual features, which is of great significance to the development of marine resources. Summary of the Invention

[0004] The purpose of the present invention is to provide an underwater target detection method based on masked self-supervised contrastive learning. This method can effectively improve the visual representation learned by the network, adapt to complex and changing underwater scenes, and improve the performance of the detector. To achieve the above purpose, the present invention provides the following technical solution, which includes the following steps:

[0005] Step S1: Construct an underwater target detection benchmark dataset.

[0006] Step S2: Constructing a dataset for mask self-supervised contrastive learning. The underwater target detection benchmark dataset from step S1 is cropped according to the target boxes in the annotation file. The cropped images are saved and classified as the dataset for the mask self-supervised contrastive learning process.

[0007] Step S3: Build a mask self-supervised contrastive learning framework. This framework consists of two branches: an online branch and a target branch. Both branches have a similar architecture: a Swin-Transformer encoder, an Adaptive Feature Fusion Pyramid Module (AFFPM), and a linear projection layer. To ensure an asymmetric structure between the two branches, an additional prediction head is introduced in the online branch.

[0008] Furthermore, the online branch updates parameters via gradients, while the target branch updates parameters via exponentially weighted average (EMA).

[0009] Furthermore, in order to balance speed and accuracy, Swin-Base in Swin-Transformer is selected as the encoder.

[0010] Furthermore, AFFPM uses a weighted fusion structure by setting learnable weights, allowing the network to determine the importance of features at different levels, thereby improving the quality of feature fusion. The fused output is expressed as follows:

[0011] Out i =w i3 P3+w i4 P4+w i5 P5+w i6 P6++w i7 P7

[0012] Where i represents the different levels of the pyramid, i∈{3, 4, 5, 6, 7}. The variable w is a learnable parameter used to fuse features from different levels.

[0013] Furthermore, the linear projection layer is composed of a 2-layer multi-layer perceptron (MLP). The prediction head has the same structure as the linear projection layer, and the loss function is defined as:

[0014]

[0015] Here q represents online view, k + is the target feature of other views of the same image, k i is the key-value queue of the target feature, τ is the temperature coefficient, and K is the length of the queue, that is, the batch size.

[0016] Furthermore, to prevent excessive mask blocks from causing large differences in views between the online branch and the target branch, which in turn affects semantic consistency in contrastive learning, the present invention only adds 25% of the mask blocks. The parameters are continuously updated through contrastive loss between different views, and the weights of Swin-Base and AFFPM are preserved.

[0017] Step S4: Fine-tune the adaptive feature fusion pyramid network AFFPM. The fine-tuning framework includes Swin-Base, AFFPM, and Cls-head (classification head).

[0018] Furthermore, the classification head pools all the hierarchical features of AFFPM, and then outputs classification features of different levels through a fully connected layer. Finally, these classification features are fused by setting the learnable weight W.

[0019] After completing the fine-tuning framework, freeze the Swin-Base network. That is, during training, do not update its parameters. Update the AFFPM network parameters. Use the common cross-entropy loss as the loss function. Input the dataset from step S2 and train. After training, save the weights of Swin-Base and AFFPM.

[0020] Step S5: Build the underwater target detection framework. Improve the YOLOX detection algorithm by replacing the YOLOX backbone network with Swin-Base and the pyramid network with AFFPM. The improved YOLOX consists of Swin-Base, AFFPM, and the Decoupled Detection Head.

[0021] Furthermore, the underwater object detection framework is initialized using the weights in step S4. The loss function includes category loss, bounding box loss, and IOU loss. It can be expressed as:

[0022]

[0023] Here λ is the balance coefficient, and is the true annotation box g i and the predicted box p j The classification loss and regression loss between . It is the IOU loss between the real annotation box and the predicted box. The classification loss uses focal loss, the regression loss uses GIOU loss, and the IOU loss uses smooth L1 loss. These loss functions are expressed as follows:

[0024]

[0025] Where α is the modulation factor, p is the predicted value, γ is a constant, and y is the label value. c is the area of the minimum bounding rectangle between the predicted box and the true box. x is the difference between the predicted value and the true value.

[0026] After the framework is built, the weights from step S4 are loaded to initialize the parameters of Swin-Base and AFFPM. The dataset from step S1 is then used for training. After training, actual underwater data images are used for inference to determine the location and category of the final target.

[0027] The beneficial effects that can be achieved by applying the above steps are:

[0028] (1) By introducing a self-supervised learning method, pre-training on ImageNet and fine-tuning on underwater datasets, it can better learn visual representations and adapt to complex and changing underwater environments;

[0029] (2) To ensure the consistency of upstream and downstream tasks, an adaptive feature fusion pyramid module is added to mask self-supervised contrastive learning, focusing on the extraction of multi-scale features;

[0030] (3) The fine-tuning process introduces a weighted classification head to process multi-scale features, further optimizing multi-scale feature extraction;

[0031] (4) By introducing mask blocks in contrastive learning, more spatial information is added, combining the advantages of mask image modeling (MIM) method and instance discrimination (ID method). BRIEF DESCRIPTION OF THE DRAWINGS

[0032] Figure 1 This is a structural diagram of the masked self-supervised contrastive learning framework of the present invention;

[0033] Figure 2 This is a structural diagram of the fine-tuning framework of the present invention;

[0034] Figure 3 This is a structural diagram of the underwater target detection framework of the present invention;

[0035] Figure 4 This is a structural diagram of the classification head of the present invention;

[0036] Figure 5 This is a structural diagram of the adaptive feature fusion pyramid module of the present invention;

[0037] Figure 6 Flowchart of the present invention. DETAILED DESCRIPTION

[0038] The following is a clear and complete description of the technical solutions in the embodiments of the present invention in conjunction with the accompanying drawings. The present invention relates to an underwater target detection method based on mask self-supervised contrastive learning, which specifically includes the following steps:

[0039] Step S1: Construct an underwater target detection benchmark dataset.

[0040] Step S2: Constructing a dataset for mask self-supervised contrastive learning. The underwater target detection benchmark dataset from step S1 is cropped according to the target boxes in the annotation file. The cropped images are saved and classified as the dataset for the mask self-supervised contrastive learning process.

[0041] Step S3: Build a mask self-supervised contrastive learning framework. Figure 1 As shown in Figure 1, the framework consists of two branches: an online branch and a target branch. Both branches have a similar architecture: a Swin-Transformer encoder, an Adaptive Feature Fusion Pyramid Module (AFFPM), and a linear projection layer. To ensure an asymmetric structure between the two branches, the online branch introduces an additional prediction head. The online branch updates parameters using gradients, while the target branch updates parameters using an exponentially weighted average (EMA).

[0042] Swin-Transformer has a layered architecture that allows cross-window connections while setting shift windows. It limits self-attention to a local window, reducing computational complexity and improving efficiency. In order to balance speed and accuracy, Swin-Base in Swin-Transformer is selected as the encoder. AFFPM adopts a weighted fusion structure, such as Figure 5 As shown in Figure 2. S3, S4, and S5 are the outputs from the last three layers of the Swin-Transformer. P3, P4, P5, P6, and P7 are similar to the structure of a traditional FPN (feature pyramid network). The variable w is a learnable parameter used to fuse features from different levels. This allows the network to independently determine the importance of features at different levels, improving the quality of feature fusion. The fused output is expressed as follows:

[0043] Out i =w i3 P3+w i4 P4+w i5 P5+w i6 P6++w i7 P7

[0044] Where i represents the different levels of the pyramid, i∈{3, 4, 5, 6, 7}. The linear projection layer consists of a 2-layer multi-layer perceptron (MLP). The prediction head has the same structure as the linear projection layer. The loss function is defined as:

[0045]

[0046] Here q represents online view, k +is the target feature of other views of the same image, k i is the key-value queue of the target feature, τ is the temperature coefficient, and K is the length of the queue, that is, the batch size.

[0047] After the framework is built, the ImageNet dataset is input. Using random cropping, different views are extracted and then scaled to the same size, for example, 384×384, for training. The views of the online branch are called online views, and mask blocks are only added to the online views. To prevent excessive mask blocks from causing large differences in views between the online and target branches, which would affect semantic consistency in contrastive learning, the present invention only adds 25% of the mask blocks. Using the contrastive loss between different views, the parameters are continuously updated, and the weights of Swin-Base and AFFPM are preserved.

[0048] Step S4: Fine-tune the adaptive feature fusion pyramid network AFFPM. The fine-tuning framework is as follows Figure 2 As shown, it includes Swin-Base, AFFPM and Cls-head (classification header). The classification header structure is as follows Figure 4 As shown in Figure 1, it pools all the hierarchical features of AFFPM, then outputs the classification features of different levels through a fully connected layer. Finally, these classification features are fused by setting the learnable weight W.

[0049] After setting up the fine-tuning framework, load the Swin-Base and AFFPM weights from step S3 as initialization. Freeze the Swin-Base network, meaning its parameters will not be updated during training. Update the AFFPM network parameters instead. Use the standard cross-entropy loss as the loss function. Enter the dataset from step S2 and train. After training, save the Swin-Base and AFFPM weights.

[0050] Step S5: Build an underwater target detection framework. Replace the backbone network in YOLOX with Swin-Base and the pyramid network with AFFPM to improve the YOLOX detection algorithm. Figure 3 As shown in Figure 1, the improved YOLOX consists of Swin-Base, AFFPM, and Decoupled Detection Head. The Decoupled Detection Head separates the features of the classification and regression tasks, handles the coupling between the classification and regression tasks, and further improves the detection accuracy. The underwater object detection framework uses the weights in step S4 to initialize the improved YOLOX. The loss function includes category loss, bounding box loss, and IOU loss. It can be expressed as:

[0051]

[0052] Here λ is the balance coefficient, and is the true annotation box g i and the predicted box p j The classification loss and regression loss between . It is the IOU loss between the real annotation box and the predicted box. The classification loss uses focal loss, the regression loss uses GIOU loss, and the IOU loss uses smooth L1 loss. These loss functions are expressed as follows:

[0053]

[0054] Where α is the modulation factor, p is the predicted value, γ is a constant, and y is the label value. c is the area of the minimum bounding rectangle between the predicted box and the true box. x is the difference between the predicted value and the true value.

[0055] After the framework is built, the weights from step S4 are loaded to initialize the parameters of Swin-Base and AFFPM. The dataset from step S1 is then used for training. After training, actual underwater data images are used for inference to determine the location and category of the final target.

Claims

1. A method for underwater target detection based on masked self-supervised contrastive learning, characterized by: The following steps are involved: Step S1: Construct an underwater target detection benchmark dataset; Step S2: Construct a mask self-supervised contrastive learning dataset; Step S3: Build a mask self-supervised contrastive learning framework, input the ImageNet dataset for pre-training, and save the weights after pre-training is completed; Step S4: Fine-tune the adaptive feature fusion pyramid network AFFPM. The fine-tuning framework includes Swin-Base, AFFPM, and Cls-head (classification head). After fine-tuning, save the weights. Step S5: Build an underwater target detection framework and perform inference after training. The contrastive learning framework in step S3 includes an online branch and a target branch; both the online branch and the target branch include a Swin-Transformer encoder, an AFFPM (Adaptive Feature Fusion Pyramid Module) adaptive feature fusion pyramid module, and a linear projection layer; the online branch also introduces an additional prediction head; the online branch updates parameters through gradients, while the target branch updates parameters through exponentially weighted average (EMA).

2. The underwater target detection method based on mask self-supervised contrastive learning according to claim 1, characterized in that: The specific method of step S2 is: cropping the underwater target detection benchmark dataset constructed in step S1 according to the target box in the annotation file, saving the cropped images and classifying them as a mask self-supervised contrast learning dataset.

3. The underwater target detection method based on mask self-supervised contrastive learning according to claim 1, characterized in that: The Swin-Transformer encoder has a layered architecture that allows cross-window connections while setting a shift window. The AFFPM adopts a weighted fusion structure, and the fused output is expressed as follows: Out i =in i3 P3+in i4 P4+w i5 P5+w i6 P6++w i7 P7 Among them, P3, P4, P5, P6, and P7 represent the structure in the traditional FPN (feature pyramid network). The variable w is a learnable parameter used to fuse features at different levels. i represents different levels of the pyramid, i∈{3, 4, 5, 6, 7}. The linear projection layer consists of a 2-layer multi-layer perceptron (MLP). The prediction head has the same structure as the linear projection layer. The loss function is defined as: Here q represents online view, k + is the target feature of other views of the same image, k i is the key-value queue of the target feature, τ is the temperature coefficient, and K is the length of the queue, that is, the batch size.

4. The underwater target detection method based on mask self-supervised contrastive learning according to claim 1, characterized in that: After the framework is built in step S3, the ImageNet dataset is input, and different views are extracted by random cropping, and then scaled to the same size for training. The views of the online branch are called online views. Mask blocks are only added to the online views. In order to prevent too many mask blocks from causing excessive differences in views between the online branch and the target branch, thereby affecting the semantic consistency in contrastive learning, only 25% of the mask blocks are added.

5. The underwater target detection method based on mask self-supervised contrastive learning according to claim 1, characterized in that: In step S4, Cls-head pools all hierarchical features of AFFPM, then outputs classification features of different levels through a fully connected layer, and finally fuses these classification features by setting a learnable weight W.

6. The underwater target detection method based on mask self-supervised contrastive learning according to claim 1, characterized in that: In step S4, after completing the fine-tuning framework, the weights of Swin-Base and AFFPM in step S3 are loaded as initialization and the Swin-Base network is frozen.

7. The underwater target detection method based on mask self-supervised contrastive learning according to claim 1, characterized in that: Step S5 replaces the backbone network in YOLOX with Swin-Base and the pyramid network with AFFPM to improve the YOLOX detection algorithm. The improved YOLOX consists of Swin-Base, AFFPM and Decoupled Detection Head. The underwater target detection framework uses the weights in step S4 to initialize the improved YOLOX. The loss function includes category loss, bounding box loss and IOU loss, which can be expressed as follows: where λ is the balance coefficient, and is the true annotation box g i and the predicted box p j The classification loss and regression loss between It is the IOU loss between the real annotation box and the predicted box, the classification loss uses focal loss, the regression loss GIOU loss, and the IOU loss uses smooth L1 loss. These loss functions are expressed as follows: Where α is the modulation factor, p is the predicted value, γ is a constant, y is the label value, and A c is the area of the minimum enclosing rectangle between the predicted box and the true box, and x is the difference between the predicted value and the true value; After the framework is built, the weights in step S4 are loaded to initialize the parameters of Swin-Base and AFFPM. Then, the dataset in step S1 is input for training. After training, the actual underwater data image is input for inference to obtain the final location and category of the target to be detected.

Images