Lightweight warship vital point recognition method based on improved YOLOv5

Abstract

The application discloses a kind of light-weight warship vital point identification methods based on improved YOLOv5, including the following steps: making warship vital point dataset-improve YOLOv5 model: last layer C3 is replaced with C3STR module in the Backbone part of YOLOv5 model;In the part of multi-scale fusion of Neck, C3 module is replaced with C3EE module, and a light attention Shuffle Attention module is added after each C3EE layer, and GSConv is replaced with Conv in the part of PAN feature fusion;In the part of Prediction Head, two scale detection branches are added-test.The application uses the light-weight warship vital point identification method based on the improved YOLOv5 described above, with the advantages of light weight, easy deployment, and can realize the identification and detection of warship vital point in real time.

Description

TECHNICAL FIELD

[0001] The application relates to the technical field of target recognition, in particular to a lightweight warship vital part recognition method based on an improved YOLOv5. BACKGROUND

[0002] In modern information-based sea battles, under different tactical battle backgrounds, the vital parts of warships have different strike values. How to effectively recognize the vital parts such as gun emplacements, radars and antennas of warships is a key link for locking and guiding the destruction of targets.

[0003] However, because the vital parts are relatively small targets, they have the characteristics of low pixel ratio and are easily blocked, and are affected by factors such as low recognition accuracy in actual recognition tasks. Therefore, improving the detection accuracy of vital parts and ensuring the detection speed of the model are the goals to be achieved in actual combat tasks.

[0004] Meanwhile, in the ImageNet competition in 2012, the AlexNet algorithm achieved outstanding results in the field of image classification, improving the classification accuracy of one thousand classes of images to 84.7%. Subsequently, the two-stage target detection algorithms represented by FasterR-CNN and MaskR-CNN, and the single-stage target detection algorithms represented by SSD and YOLO, have achieved excellent recognition performance on large-scale target detection data sets.

[0005] Especially in 2015, YOLOv1 was proposed, and its core idea is similar to FasterR-CNN. It takes the entire image as the input of the network and directly determines the position and classification of the bounding box in the output layer. Further, YOLOv2 adopts batch normalization, high-class classifier and prior box optimization strategies, and realizes the improvement of speed, accuracy and the number of recognized object categories. Further, the structure of YOLOv3 is more complex, it introduces the feature pyramid network (FPN) and darknet-53 network, and allows the network structure to be changed to balance speed and accuracy, thereby greatly improving the calculation speed. Further, the network architecture of YOLOv4 changes a lot, through a large number of parameter tuning experiments, the best balance between input network resolution, convolution layer number and parameter number is found, and the comprehensive performance is improved. Further, YOLOv5 has obvious contribution in network lightweight, faster speed and easier deployment.

[0006] Correspondingly, the prior art adopts the following processing methods in the research of ship detection: 1. Based on the segmentation algorithm, the connected regions in the image are effectively utilized, the histogram of oriented gridients (HOG) features are extracted to train the classifier, and the offline classification and online detection of the ship target are realized; 2. The Faster R-CNN algorithm is combined with the constant false alarm rate, the object suggestions generated by the algorithm are used as the protection window, after the small size target is obtained, the prediction box with relatively low score is re-evaluated, and the ship target detection effect in the complex background is improved; 3. The attention mechanism is introduced into RetinaNet, the FocalLoss is used to improve the sample imbalance and other problems, and in the ship detection experiment of multi-resolution imaging, high detection accuracy is achieved; 4. The serial correction linear unit is added to the shallow network of SSD, the feature transmission efficiency is improved, the feature pyramid network is used to gradually fuse the multi-scale feature maps of different levels, and the positioning and classification accuracy of the ship target is improved;

[0007] It can be known that the above methods all lack the detection method of distinguishing the vital parts of the ship, and the speed and accuracy of the detection cannot meet the real-time requirements of the battlefield. In the actual captured ship image, the proportion of vital targets such as gun emplacements, radars and antennas is small and easy to be blocked, and the proportion of false detection and missed detection is significantly increased. SUMMARY

[0008] To solve the above problems, the present application provides a lightweight ship vital part recognition method based on improved YOLOv5, which increases two target detection branches, and takes into account the detection capability of the model for large ship targets and small vital part targets, especially for the situation that the vital part positions of the ship are relatively dense and easy to be blocked, and can show better detection effect, and has high accuracy for small vital part target recognition. At the same time, the basic model adopts Yolov5 which has detection accuracy and speed, has the advantages of lightweight and easy deployment, effectively improves the detection capability under the premise of increasing a small amount of parameters, can be deployed on the missile-borne equipment with limited computing power, realizes the recognition and detection of the vital parts of the ship in real time, thereby solving the problems of false detection and missed detection of the original model for vital parts, and improving the detection accuracy of the model for large ship targets and small vital part targets.

[0009] To achieve the above purpose, the present application provides a lightweight ship vital part recognition method based on improved YOLOv5, comprising the following steps:

[0010] S1, collect the original ship image, and make a ship vital part data set;

[0011] S2, improve the YOLOv5 model:

[0012] Two layers of down-sampling layers are added to the Backbone part of the YOLOv5 model, a Swin Transformer module is introduced, and the last C3 layer is replaced with a C3STR module;

[0013] In the Neck multi-scale fusion part, two scale fusion layers are added, a C3EE module is used to replace the C3 module, and a lightweight attention Shuffle Attention module is added after each C3EE layer. In the PAN feature fusion part, a GSConv is used to replace the Conv;

[0014] In the Prediction Head part, two scale detection branches are added;

[0015] S3, divide the warship vital part dataset into a training set and a test set, and input the training set into the improved YOLOv5 model for model training. After the training is completed, the test set is input into the improved YOLOv5 model for warship vital part testing.

[0016] Preferably, step S1 specifically comprises the following steps:

[0017] Online search image materials with warship keywords using online crawler technology as warship original images, and label them according to four categories of warship, fort, radar and antenna through Labelme label making software, and then convert them into txt data format required by Yolo training to complete the warship vital part dataset.

[0018] Preferably, the Backbone part of step S2 comprises 2 times down-sampling layer, 4 times down-sampling layer, C3 module, 8 times down-sampling layer, C3 module, 16 times down-sampling layer, C3 module, 32 times down-sampling layer, C3 module, 64 times down-sampling layer, C3 module, 128 times down-sampling layer, C3STR module and SPPF module connected in sequence, that is, 32 times down-sampling layer and 128 times down-sampling layer are added;

[0019] It performs the following operations:

[0020] The input warship original image with a size of 640*640*3 is subjected to slicing operation through a 2-fold downsampling layer to reduce the image size and increase the channel number, so that the feature map size becomes 320*320*64; the feature map size becomes 160*160*128 through a 4-fold downsampling layer; the feature map size becomes 80*80*256 through an 8-fold downsampling layer; the feature map size becomes 40*40*512 through a 16-fold downsampling layer; the feature map size becomes 20*20*768 through a 32-fold downsampling layer; the feature map size becomes 10*10*1024 through a 64-fold downsampling layer; and the feature map size becomes 5*5*1280 through a 128-fold downsampling layer.

[0021] Preferably, the traversal number of the Bottleneck and the STRBottleneck in the C3 module and the C3STR module is 3.

[0022] Preferably, in the C3STR module, a Swin Transformer feature extraction backbone network is introduced, which uses a multi-head attention mechanism Windows Multi-HeadSelf-Attention to divide the extracted feature map into multiple disjoint regions, and the W-MSA attention mechanism only calculates within each window.

[0023] Preferably, the Neck multi-scale fusion part in step S2 includes 8-scale fusion layers, 16-scale fusion layers, 32-scale fusion layers, 64-scale fusion layers and 128-scale fusion layers arranged corresponding to the output feature maps of the Backbone part, that is, 64-scale fusion layers and 128-scale fusion layers are added.

[0024] The C3EE module includes four input branches, and the channel number of each input branch is 1 / 4 of the original channel number; two 1*1 convolution kernels and three 3*3 convolution kernels introduced through the two added input branches form a multi-branch structure.

[0025] In the Shuffle Attention module, the features are divided into multiple sub-features according to the channel, and the channel and spatial attention operations are performed on each sub-feature; a shuffle unit is used to combine and process the channel and spatial attention modules of the above sub-features to perform adaptive feature refinement.

[0026] The GSConv first halves the input feature channels, then generates another half of the channel feature map through a depth separation convolution, and then performs Concate splicing to output a reshape-adjusted feature map through a Shuffle network.

[0027] Preferably, the Shuffle Attention module will input features X∈R C×H×W Divide the channels into G groups X = [X1, X2, ..., X...]. G ],X k ∈R C / G×H×W Each feature group acquires a specific semantic response during training. At the start of Attention, for each group of X... k Both are divided into two branches X k1 ,X k2 ∈R C / 2G×H×W , where X k1 Channel attention is generated by utilizing the interrelationships between channels, X k2 Spatial attention is generated by utilizing the spatial relationships of features. The attention module generates an importance coefficient for each sub-feature, which is then multiplied with the original feature map.

[0028] Preferably, the Neck multi-scale fusion part in step S2 is an FPN+PAN network structure, that is, multi-scale fusion is performed using bottom-up and top-down pathways, and the specific steps are as follows:

[0029] The feature maps with dimensions of 80×80×256, 40×40×512, 20×20×768, 10×10×1024, and 5×5×1280, respectively, are input from the Backbone part and are denoted as M1, M2, M3, M4, and M5, respectively.

[0030] First, the feature map is processed through an FPN structure, which performs a bottom-up multi-scale feature fusion to integrate deep semantic information with shallow detail information: M5 is processed by a 1×1 convolution to eliminate the aliasing effect caused by fusion, and is denoted as P5; P5 is upsampled by 2x and concatenated with the M4 feature map, and the resulting feature map is generated by the C3EE module and is denoted as P4; P4 is processed by a 1×1 convolution to eliminate the aliasing effect caused by fusion, and after being upsampled by 2x, it is concatenated with the M3 feature map, and the resulting feature map is generated by the C3EE module and is denoted as P3; P3 is upsampled by 2x and concatenated with the M2 feature map, and the resulting feature map is generated by the C3EE module and is denoted as P2; P2 is upsampled by 2x and concatenated with the M1 feature map, and the resulting feature map is generated by the C3EE module and is denoted as Q1.

[0031] Then the PAN structure is adopted to perform one-time top-down multi-scale fusion to form the final output feature map: Q1 is subjected to twice down-sampling and Concate splicing with the P2 feature map, a feature map recorded as Q2 is generated through the C3EE+Shuffle Attention module; Q2 is subjected to twice down-sampling and Concate splicing with the P3 feature map, a feature map recorded as Q3 is generated through the C3EE+Shuffle Attention module; Q3 is subjected to twice down-sampling and Concate splicing with the P4 feature map, a feature map recorded as Q4 is generated through the C3EE+Shuffle Attention module; Q4 is subjected to twice down-sampling and Concate splicing with the P5 feature map, a feature map recorded as Q5 is generated through the C3EE+Shuffle Attention module.

[0032] The final output feature map of the Neck part is Q1 to Q5, and the sizes are: 80x80x256, 40x40x512, 20x20x768, 10x10x1024 and 5x5x1280.

[0033] Preferably, the Prediction Head part in step S2 includes 80x80x256, 40x40x512, 20x20x768, 10x10x1024 and 5x5x1280 scale prediction branches arranged corresponding to the output of the multi-scale fusion part of the Neck, that is, 10x10x1024 and 5x5x1280 scale prediction branches are added.

[0034] The Anchors of the prediction network are set according to the clustering of the data set, 3 anchors are predicted for each scale, the output candidate frame is mapped to the original frame size through non-maximum suppression, and the target object in the image is marked to obtain the final result.

[0035] The present application has the following beneficial effects:

[0036] I. The Backbone part uses a C3STR module to replace the last C3 module. The module uses the Windows Multi-Head Self-Attention multi-head attention mechanism of the SwinTransformer network structure, divides the extracted feature map into multiple disjoint regions (Window), and the W-MSA attention mechanism only calculates within each window, ensuring information transmission while reducing computational complexity, allowing the network to extract more detailed information and enhance feature extraction capability.

[0037] II. Neck part, all C3 modules are replaced with C3EE modules. The module is inspired by the idea of Inception network, two branches are added to the C3 module, two 1x1 and three 3x3 convolution kernel convolution layers are introduced to form a multi-branch structure, two dilated convolution layers are introduced, and the rate represents the parameter of the dilated convolution layer, C3EE structure uses two diated convolution parameters of rate = 3 and rate = 5 respectively, finally the convolution layers of different sizes and rate are spliced and fused, which effectively expands the receptive field of the model;

[0038] GSConv is used to replace the traditional Conv convolution in the PAN feature fusion part, GSConv can preserve the hidden connection between channels, reduce the loss of feature information, and use GSConv to process the feature map, which has less redundant information and does not need to be compressed, and the attention module works better;

[0039] A lightweight attention mechanism Shuffle Attention module is added after each C3EE module, Shuffle Attention module fully draws on the idea of ShuffleNetv2 network, divides multiple sub-features, each sub-feature infers feature weights through channel and space dimensions, multiplies with input features, performs adaptive feature refinement, and retains more effective features into the following network structure.

[0040] III. Prediction Head part, on the basis of the original three scale prediction branch, two prediction branches of [10x10x1024] and [5x5x1280] are added, which improves the detection accuracy of the model for large ship and small target critical parts, and improves the detection accuracy of the model for large ship and small target critical parts.

[0041] The technical solutions of the present application will be further described in detail below with the help of the drawings and examples. DESCRIPTION OF DRAWINGS

[0042] Figure 1 It is the existing YOLOv5 network structure diagram.

[0043] Figure 2 It is the YOLOv5 network structure diagram of the improved YOLOv5-based lightweight warship critical part identification method.

[0044] Figure 3 It is the warship and critical part label illustration diagram labeled by Labelme software of the improved YOLOv5-based lightweight warship critical part identification method.

[0045] Figure 4The C3STR module introducing the Swin Transformer block for the improved YOLOv5-based lightweight warship vital point identification method described in the application.

[0046] Figure 5 The C3EE module structure diagram for the improved YOLOv5-based lightweight warship vital point identification method described in the application.

[0047] Figure 6 The GSConv module structure diagram for the improved YOLOv5-based lightweight warship vital point identification method described in the application.

[0048] Figure 7 The Shuffle Attention module structure diagram for the improved YOLOv5-based lightweight warship vital point identification method described in the application.

[0049] Figure 8 The bounding box loss result graph of the training set of the test example of the application.

[0050] Figure 9 The target detection loss mean result graph of the training set of the test example of the application.

[0051] Figure 10 The classification loss mean result graph of the training set of the test example of the application.

[0052] Figure 11 The bounding box loss result graph of the validation set of the test example of the application.

[0053] Figure 12 The target detection loss mean result graph of the validation set of the test example of the application.

[0054] Figure 13 The classification loss mean result graph of the validation set of the test example of the application.

[0055] Figure 14 The precision change curve graph of the test example of the application.

[0056] Figure 15 The recall rate change curve graph of the test example of the application.

[0057] Figure 16 The average map change graph of the test example of the application with a threshold greater than 0.5.

[0058] Figure 17 The IOU threshold is (0.5-0.95) average map graph of the test example of the application.

[0059] Figure 18 The detection effect diagram of a warship and a vital part using a conventional Yolov5s model for a test example of the present application.

[0060] Figure 19 The detection effect diagram of a warship and a vital part using an algorithm model of the present application for a test example of the present application. DETAILED DESCRIPTION

[0061] The present application will be further described below in combination with the drawings. It should be noted that the present embodiment is based on the technical solution, and gives a detailed implementation and specific operation process, but the protection scope of the present application is not limited to the present embodiment.

[0062] The light-weight warship vital part recognition method based on the improved YOLOv5 includes the following steps:

[0063] S1, collect warship original images, and make a warship vital part data set;

[0064] Preferably, step S1 specifically includes the following steps:

[0065] Online search image materials with warship keywords using online crawler technology as warship original images, and mark them according to four categories of warship, battery, radar and antenna through Labelme label making software, and then convert them into txt data format required by Yolo training to complete the warship vital part data set.

[0066] In the present embodiment, more than 300 pictures of various warships are searched through online crawler, and are allocated to the training set and the verification set according to a 4:1 ratio. In the data set label making process, four types of recognition targets, i.e. boat, battery, radar and antenna, are marked according to rectangular frames. After marking, the data set generates json format data containing the recognition category "label" and the coordinate "points" of the rectangular frame. Since the data set used by Yolo is in txt format, the format of the data set needs to be converted to meet the format required by Yolo training.

[0067] And the input image uses Mosaic data enhancement operation to splice four pictures in a random scaling, random cropping and random arrangement manner to improve the training speed and network precision of the model; the self-adaptive scaling method is used to adjust all pictures to the specification size of 640x640x3 to improve the training speed of the Yolov5 network

[0068] S2, improve the YOLOv5 model:

[0069] Two layers of down-sampling layers are added to the Backbone part of the YOLOv5 model, a Swin Transformer module is introduced, and the last C3 layer is replaced with a C3STR module.

[0070] Preferably, the Backbone part in step S2 comprises, in sequence, a 2-fold down-sampling layer, a 4-fold down-sampling layer, a C3 module, an 8-fold down-sampling layer, a C3 module, a 16-fold down-sampling layer, a C3 module, a 32-fold down-sampling layer, a C3 module, a 64-fold down-sampling layer, a C3 module, a 128-fold down-sampling layer, a C3STR module, and an SPPF module, that is, a 32-fold down-sampling layer and a 128-fold down-sampling layer are added.

[0071] It performs the following operations:

[0072] The input ship original image with a size of 640x640x3 is subjected to slicing operation by the 2-fold down-sampling layer to reduce the image size and increase the channel number, and the feature map size becomes 320x320x64; the feature map size becomes 160x160x128 by the 4-fold down-sampling layer; the feature map size becomes 80x80x256 by the 8-fold down-sampling layer; the feature map size becomes 40x40x512 by the 16-fold down-sampling layer; the feature map size becomes 20x20x768 by the 32-fold down-sampling layer; the feature map size becomes 10x10x1024 by the 64-fold down-sampling layer; and the feature map size becomes 5x5x1280 by the 128-fold down-sampling layer.

[0073] Preferably, the Bottleneck and STRBottleneck in the C3 module and the C3STR module each have a traversal number of 3.

[0074] In the feature extraction process of the network, the shallow feature map can better retain the edge and texture features of the image, and the deep feature map can better retain the semantic information of the image. With multiple times of down-sampling feature extraction, the loss of detailed information is more serious. To make up for the loss of information caused by the two added down-sampling layers, the C3STR module combined with the Swin Transformer network structure is introduced to replace the last C3 module in the Backbone part, which reduces the loss of detailed information in the multiple down-sampling processes without increasing a large amount of calculation, and improves the recognition ability of small targets of ships and vital parts.

[0075] Preferably, in the C3STR module, the Swin Transformer feature extraction backbone network is introduced, which uses the multi-head attention mechanism Windows Multi-HeadSelf-Attention to divide the extracted feature map into multiple disjoint regions, and the W-MSA attention mechanism only calculates within each window.

[0076] In the Neck multi-scale fusion part, two scale fusion layers are added, and a C3EE module is used to replace the C3 module. The C3EE module is inspired by the Inception network structure and introduces a dilated convolution. Two branches are added to the original two branches of the C3 module to reduce the loss of feature information in the process of multiple convolution transformations, effectively expand the receptive field of the model, and a lightweight attention Shuffle Attention module is added after each C3EE layer.

[0077] Because the spatial information is gradually transmitted to the channel, the compression and expansion of the channel will cause the loss of part of the semantic information. Dense convolution calculation maximizes the hidden connection between each channel, while sparse convolution cuts off these connections. Therefore, in the PAN feature fusion part, the GSConv is used to replace the Conv to reduce the loss of feature information caused by repeated channel compression and expansion, and to maximize the hidden connection between channels.

[0078] Preferably, the Neck multi-scale fusion part of step S2 comprises an 8-scale fusion layer, a 16-scale fusion layer, a 32-scale fusion layer, a 64-scale fusion layer, and a 128-scale fusion layer arranged corresponding to the output feature map of the Backbone part, i.e., the 64-scale fusion layer and the 128-scale fusion layer are added, so that the bottom layer semantic information applied to the feature fusion part is more abundant.

[0079] The C3EE module comprises four input branches, and the number of channels of each input branch is reduced to 1 / 4 of the original number. Two 1x1 convolution kernels and three 3x3 convolution kernels are introduced in the two added input branches to form a multi-branch structure. That is, the idea of the Inception network structure is used, and two dilated convolution layers are introduced. The parameter of the dilated convolution layer is represented by rate. In the Deeplab segmentation algorithm applied before the dilated convolution layer, the main function is to expand the receptive field. In the C3EE structure, two diated convolution parameters of rate=3 and rate=5 are used. Finally, the outputs of the convolution layers with different sizes and rates are fused to enhance the feature extraction capability of the network, fully utilize the deep and shallow layer feature information extracted by the network, and effectively expand the receptive field.

[0080] In the Shuffle Attention module, the features are divided into multiple sub-features according to the channel, and the channel and spatial attention operations are performed on each sub-feature. A shuffleunit is used to combine and process the channel and spatial attention modules of the above sub-features to perform adaptive feature refinement.

[0081] GSConv first halves the input feature channel, then generates another half channel feature map through depth separation convolution, and then performs Concate splicing, and outputs the reshape adjusted feature map through a Shuffle network. However, if GSConv is used in all stages of the model, the network layer of the model will be deeper, and the deep layer will increase the resistance to data flow, significantly increasing the inference time. When these feature maps go to the Neck part, the channel dimension reaches the maximum, and the width and height dimensions reach the minimum, and no longer need to be transformed. Therefore, only GSConv is used in the Neck part, and in this stage, the feature map processed by GSConv has less redundant information and does not need to be compressed, and the attention module works better.

[0082] Preferably, the Shuffle Attention module will input feature X∈R C×H×W According to the channel, the input feature X is divided into G groups X = [X1, X2,.., X G ], X k ∈R C / G×H×W Each group of features obtains a specific semantic response in the training process, and at the beginning of Attention, each X k in each group is divided into two branches X k1 , X k2 ∈R C / 2G×H×W , wherein X k1 generates channel attention using the interrelation between channels, and X k2 generates spatial attention using the spatial relationship of the features, and generates a corresponding importance coefficient for each sub-feature through the attention module, and then multiplies it with the original feature map, so that the model can strengthen the use of effective information.

[0083] Under the visual conditions of real-time monitoring, the small target image of the vital part on the warship generally has only a few pixels, and the network model has very little semantic information when extracting small targets. For the detection of vital parts, more top feature maps containing rich detailed information are needed. Therefore, the Neck multi-scale fusion part described in step S2 is a FPN+PAN network structure, that is, a bottom-up and top-down path is used for multi-scale fusion, so that the deep features containing rich semantic information are fused with the shallow features containing edge, texture and other detailed information. That is, the feature pyramid network transmits semantic information from top to bottom, the path aggregation network transmits texture, position and other detailed information from bottom to top, and multi-scale feature fusion is beneficial to better learning of feature information by the network and enhances the detection ability of the model for dense small targets and occluded targets. The specific execution steps are as follows:

[0084] The feature maps with sizes of 80x80x256, 40x40x512, 20x20x768, 10x10x1024, 5x5x1280 input from the Backbone part are respectively denoted as M1, M2, M3, M4, M5;

[0085] Firstly, the FPN structure is passed through, and the deep semantic information and the shallow detail information are fused through one bottom-up multi-scale feature fusion: M5 is subjected to 1x1 convolution to eliminate the aliasing effect brought by fusion, and is denoted as P5; P5 is subjected to two times of upsampling, is concatenated with the M4 feature map, and generates a feature map through a C3EE module, which is denoted as P4; P4 is subjected to 1x1 convolution to eliminate the aliasing effect brought by fusion, is subjected to two times of upsampling, is concatenated with the M3 feature map, and generates a feature map through a C3EE module, which is denoted as P3; P3 is subjected to two times of upsampling, is concatenated with the M2 feature map, and generates a feature map through a C3EE module, which is denoted as P2; P2 is subjected to two times of upsampling, is concatenated with the M1 feature map, and generates a feature map through a C3EE module, which is denoted as Q1;

[0086] Then, the PAN structure is adopted, and one top-down multi-scale fusion is performed to form the final output feature map: Q1 is subjected to two times of downsampling, is concatenated with the P2 feature map, and generates a feature map through a C3EE+Shuffle Attention module, which is denoted as Q2; Q2 is subjected to two times of downsampling, is concatenated with the P3 feature map, and generates a feature map through a C3EE+Shuffle Attention module, which is denoted as Q3; Q3 is subjected to two times of downsampling, is concatenated with the P4 feature map, and generates a feature map through a C3EE+Shuffle Attention module, which is denoted as Q4; Q4 is subjected to two times of downsampling, is concatenated with the P5 feature map, and generates a feature map through a C3EE+Shuffle Attention module, which is denoted as Q5;

[0087] The final output feature maps of the Neck part are Q1 to Q5, and the sizes are respectively: 80x80x256, 40x40x512, 20x20x768, 10x10x1024, 5x5x1280.

[0088] In the Prediction Head part, two scale detection branches are added;

[0089] Preferably, the Prediction Head part in step S2 includes 80x80x256 scale prediction branches, 40x40x512 scale prediction branches, 20x20x768 scale prediction branches, 10x10x1024 scale prediction branches, and 5x5x1280 scale prediction branches arranged corresponding to the output of the Neck multi-scale fusion part, i.e., 10x10x1024 scale prediction branches and 5x5x1280 scale prediction branches are added, so as to fully utilize the feature maps of various scales extracted by the network and improve the prediction accuracy.

[0090] The Anchors of the prediction network are set according to the clustering of the data set in advance, 3 anchors are predicted for each scale, the output candidate frame is mapped to the original frame size through non-maximum suppression, and the target object in the image is marked to obtain the final result.

[0091] The setting of the preselected frame size greatly affects the detection accuracy and convergence speed of the model. The default size of the Anchors in Yolov5 is set according to the COCO data set. For the warship vital part data set, there are both large warship targets and small vital part targets with relatively dense positions, so it is necessary to reset the anchor size. The YOLOv5 automatic anchor calculation method is used in this embodiment, the K-means algorithm is used for adaptive anchor calculation, and for the five output feature map sizes, 3 Anchors with different width-height ratios are generated for each size, which are [8, 8], [12, 12], [15, 14], [18, 18], [26, 23], [41, 19], [40, 31], [59, 40], [64, 90], [204, 144], [541, 163], [350, 299], [568, 242], [509, 349], and [424, 488].

[0092] S3, the warship vital part data set is divided into a training set and a test set, and the training set is input into the improved YOLOv5 model for model training. After training, the test set is input into the improved YOLOv5 model for warship vital part testing.

[0093] Test example

[0094] The model weight is initialized in the experiment, the optimizer is Adam under the keras framework, and the hyperparameters are set as follows: the initial learning rate Lr0 is set to 0.01, the training batch size batch-size is set to 4, and the iteration number epochs is set to 100. Figure 10As shown, after 50 epochs of iteration, the loss is reduced to below 0.06; after 100 epochs of iteration, the loss is reduced to below 0.05; therefore, the model after 100 batches of iteration is taken for experimental testing in this embodiment.

[0095] In this test example, the precision P, recall R and AP value are taken as the model evaluation indexes, and the formulas are as follows:

[0096]

[0097]

[0098] AP = ∫P d (R)

[0099] wherein, T P represents the number of true positives, F P represents the number of false positives, and F N represents the number of false negatives.

[0100] The prediction box confidence threshold conf-thres is set to 0.25, and the non-maximum suppression iou-thres is set to 0.45, the corresponding recall R, accuracy P and mAP average precision are calculated, mAP0.5 is the average precision when the Iou threshold is set to 0.5, and the test results of the item-by-item ablation experiment on the improved YOLOv5 module are as follows:

[0101] Table 1 is the test result table

[0102]

[0103] As shown in Table 1, the model mAP@0.5 is improved, according to the above table, the mAP@0.5 value of the improved C3EE structure based on Yolov5 is improved overall, the Ture model of C3EE has a shotcut residual, the mAP@0.5 of all recognition categories allclass is improved by 0.039, the False model of C3EE has no shotcut residual, the mAP@0.5 of all recognition categories allclass is improved by 0.027; the mAP@0.5 of all recognition categories allclass is improved by 0.022 when the improved C3STR module is used in the Backbone; the mAP@0.5 of all recognition categories allclass is improved by 0.041 when the GS+ShuffleAttention module is added in the Neck part; after the above three improvement measures, the mAP@0.5 of all recognition categories allclass is improved by 0.045. Therefore, the improvement measures taken in the experiment have achieved certain optimization results in the actual detection task.

[0104] As Figures 8-19 indicated (need to explain Figures 8-17 The horizontal coordinate in each of the above figures represents the training epoch, and the vertical coordinate represents the training set bounding box loss mean value, the training set target detection loss mean value, the training set classification loss mean value, the validation set bounding box loss mean value, the validation set target detection loss mean value, the validation set classification loss mean value, the precision change value, the recall change value, the average map change value with a threshold greater than 0.5, and the average map change value with an IOU threshold of (0.5-0.95). Compared with the original Yolov5 algorithm detection result, the improved YOLOv5-based lightweight warship vital position recognition method can more accurately detect the vital position of a warship, solve the problem of false detection and missed detection of small targets in vital parts, and improve the detection accuracy of the model for both large targets and small targets in vital parts of a warship.

[0105] Therefore, the improved YOLOv5-based lightweight warship vital position recognition method replaces the C3 module in the Backbone part with a C3STR module capable of reducing feature loss, replaces the C3 module in the Neck part with a C3EE module capable of expanding the receptive field, uses a GSConv convolution and a lightweight attention mechanism ShuffleAttention module capable of maintaining hidden connections between channels in the FPN structure, and improves the recognition and detection accuracy of small targets in vital parts of a warship.

[0106] Finally, it should be noted that the above examples are only used to illustrate the technical solutions of the present application and not to limit them. Although the present application has been described in detail with reference to the preferred embodiments, those skilled in the art should understand that the technical solutions of the present application can still be modified or replaced by equivalents, and these modifications or replacements should not make the modified technical solutions deviate from the spirit and scope of the technical solutions of the present application.

Claims

1. A method for identifying key parts of lightweight warships based on an improved YOLOv5, characterized in that: Includes the following steps: S1. Collect original images of warships and create a dataset of key parts of warships; S2. Improved YOLOv5 model: Two downsampling layers were added to the backbone of the YOLOv5 model, the Swing Transformer module was introduced, and the last layer C3 was replaced with the C3STR module. In the Neck multi-scale fusion part, two scale fusion layers are added, the C3 module is replaced with the C3EE module, and a lightweight attention Shuffle Attention module is added after each C3EE layer. In the PAN feature fusion part, Conv is replaced with GSConv. In the Prediction Head section, two detection branches at different scales are added; The Backbone section mentioned in step S2 includes a 2x downsampling layer, a 4x downsampling layer, a C3 module, an 8x downsampling layer, a C3 module, a 16x downsampling layer, a C3 module, a 32x downsampling layer, a C3 module, a 64x downsampling layer, a C3 module, a 128x downsampling layer, a C3STR module, and an SPPF module connected in sequence, that is, 32x downsampling layers and 128x downsampling layers are added; It performs the following operations: An original input warship image of size 640×640×3 is sliced ​​using a 2x downsampling layer to reduce the image size and increase the number of channels, resulting in a feature map size of 320×320×64. After a 4x downsampling layer, the feature map size becomes 160×160×128; after an 8x downsampling layer, it becomes 80×80×256; after a 16x downsampling layer, it becomes 40×40×512; after a 32x downsampling layer, it becomes 20×20×768; after a 64x downsampling layer, it becomes 10×10×1024; and after a 128x downsampling layer, it becomes 5×5×1280. The Neck multi-scale fusion part described in step S2 includes 8-scale fusion layers, 16-scale fusion layers, 32-scale fusion layers, 64-scale fusion layers, and 128-scale fusion layers arranged according to the output feature maps of the Backbone part, that is, 64-scale fusion layers and 128-scale fusion layers are added. The C3EE module includes 4 input branches, and the number of channels in each input branch is reduced to 1 / 4 of the original. The two additional input branches introduce two 1×1 convolution kernels and three 3×3 convolution kernels to form a multi-branch structure. In the Shuffle Attention module, the features are divided into multiple sub-features according to the channel segmentation features. Channel and spatial attention operations are performed on each sub-feature. A shuffle unit is used to merge and process the channel and spatial attention modules of the above sub-features to perform adaptive feature refinement. GSConv first halve the input feature channels, then generate the other half of the channel feature maps through depthwise splitting convolution, then concatenate them, and output the reshaped feature maps through a Shuffle network. S3. The dataset of key parts of the warship is divided into a training set and a test set. The training set is input into the improved YOLOv5 model for model training. After training, the test set is input into the improved YOLOv5 model for testing of key parts of the warship.

2. The method for identifying key parts of lightweight warships based on the improved YOLOv5 according to claim 1, characterized in that: Step S1 specifically includes the following steps: We used web crawling technology to search for image materials containing the keyword "warship" online, which were then used as the original images of warships. We then used Labelme tagging software to label them according to four categories: warship, gun turret, radar, and antenna. Finally, we converted them into the txt data format required for YOLO training to create a dataset of key parts of warships.

3. The method for identifying key parts of lightweight warships based on the improved YOLOv5 according to claim 1, characterized in that: In both the C3 and C3STR modules, Bottleneck and STRBottleneck are traversed 3 times.

4. The method for identifying key parts of lightweight warships based on the improved YOLOv5 according to claim 1, characterized in that: In the C3STR module, the Swing Transformer feature extraction backbone network is introduced, which uses the Windows Multi-Head Self-Attention mechanism to divide the extracted feature map into multiple non-overlapping regions. The W-MSA attention mechanism is only computed within each window.

5. The method for identifying key parts of lightweight warships based on the improved YOLOv5 according to claim 1, characterized in that: The Shuffle Attention module takes the input features X∈R as input features. C×H×W Divide the channels into G groups X = [X1, X2, ..., X...]. G ],X k ∈R C / G×H×W Each feature group acquires a specific semantic response during training. At the start of Attention, for each group of X... k Both are divided into two branches X k1 ,X k2 ∈R C / 2G×H×W , where X k1 Channel attention is generated by utilizing the interrelationships between channels, X k2 Spatial attention is generated by utilizing the spatial relationships of features. The attention module generates an importance coefficient for each sub-feature, which is then multiplied with the original feature map.

6. The method for identifying key parts of lightweight warships based on improved YOLOv5 according to claim 1, characterized in that: The Neck multi-scale fusion part mentioned in step S2 is an FPN+PAN network structure, that is, it uses bottom-up and top-down pathways to perform multi-scale fusion, and it specifically performs the following steps: The feature maps with dimensions of 80×80×256, 40×40×512, 20×20×768, 10×10×1024, and 5×5×1280, respectively, are input from the Backbone part and are denoted as M1, M2, M3, M4, and M5, respectively. First, the feature map is processed through an FPN structure, which performs a bottom-up multi-scale feature fusion to fuse deep semantic information with shallow detail information: M5 is processed by a 1×1 convolution to eliminate the aliasing effect caused by fusion, and is denoted as P5; P5 is upsampled by 2x and concatenated with the M4 feature map, and the resulting feature map is generated by the C3EE module and is denoted as P4; P4 is processed by a 1×1 convolution to eliminate the aliasing effect caused by fusion, and is upsampled by 2x and concatenated with the M3 feature map, and the resulting feature map is generated by the C3EE module and is denoted as P3; P3 is upsampled by 2x and concatenated with the M2 feature map, and the resulting feature map is generated by the C3EE module and is denoted as P2; P2 is upsampled by 2x and concatenated with the M1 feature map, and the resulting feature map is generated by the C3EE module and is denoted as Q1. Then, using a PAN structure, a top-down multi-scale fusion is performed to form the final output feature map: Q1 is downsampled by a factor of 2 and concatenated with the P2 feature map, and the resulting feature map is generated by the C3EE+Shuffle Attention module, denoted as Q2; Q2 is downsampled by a factor of 2 and concatenated with the P3 feature map, and the resulting feature map is generated by the C3EE+Shuffle Attention module, denoted as Q3; Q3 is downsampled by a factor of 2 and concatenated with the P4 feature map, and the resulting feature map is generated by the C3EE+Shuffle Attention module, denoted as Q4; Q4 is downsampled by a factor of 2 and concatenated with the P5 feature map, and the resulting feature map is generated by the C3EE+Shuffle Attention module, denoted as Q5. The final output feature maps of the Neck section are Q1 to Q5, with sizes of 80×80×256, 40×40×512, 20×20×768, 10×10×1024, and 5×5×1280, respectively.

7. The method for identifying key parts of lightweight warships based on improved YOLOv5 according to claim 1, characterized in that: The Prediction Head section mentioned in step S2 includes 80×80×256 scale prediction branches, 40×40×512 scale prediction branches, 20×20×768 scale prediction branches, 10×10×1024 scale prediction branches, and 5×5×12805 scale prediction branches corresponding to the output arrangement of the Neck multi-scale fusion section, that is, 10×10×1024 scale prediction branches and 5×5×1280 scale prediction branches are added; The prediction network predicts three anchors at each scale based on the clustering settings of the dataset. Through non-maximum suppression, the output candidate boxes are mapped to the original image size, and the target objects in the image are marked to obtain the final result.

Images