Target detection model training method and device, electronic equipment and storage medium

Abstract

The invention provides a target detection model training method and device, electronic equipment and a storage medium, and relates to the technical field of computer processing. The method comprises the following steps: acquiring a data set for model training, and inputting the data set into a feature extraction network in a to-be-trained initial detection model to perform multi-scale feature extraction and attention feature extraction on sample images to obtain a plurality of multi-scale feature maps corresponding to each sample image; inputting each multi-scale feature map into a detection head in the initial detection model to obtain a detection result corresponding to the sample image, then obtaining a loss value between each to-be-detected target and a real target corresponding to each to-be-detected target in the detection result, judging whether the loss value meets a preset model convergence condition, and if yes, determining that the to-be-detected target does not meet the preset model convergence condition; and if yes, taking the initial detection model meeting a preset model convergence condition as a target detection model. In this way, the problems that a traditional target recognition model is low in small target detection precision and poor in effect in the video satellite image can be solved.

Description

Technical Field

[0001] The present invention relates to the field of computer processing technology, and in particular to a target detection model training method, device, electronic device and storage medium. Background Art

[0002] Video satellites play an important role in monitoring, navigation, resource management and other fields. Small target detection in video satellite images is a key task in target recognition in video satellite images. As a hot and difficult problem in the field of computer vision, especially for video satellite images, due to the long shooting distance, the low resolution of the ground target itself, the target occupies fewer pixels in the image, and is interfered by complex background factors, it becomes very difficult to accurately detect tiny targets directly from the input image.

[0003] At present, network models used for small target detection usually use attention mechanisms to improve the performance of the model for small target detection, such as SE (Squeeze-and-Excitation), ECA (Efficient Channel Attention), CBAM (Convolutional Block Attention Module), SGE (Spatial Group-wise Enhance), EMA (Efficient Multi-Scale Attention), etc. The existing attention mechanism does not fully consider the level of feature extraction, and there are problems such as low accuracy and poor effect in small target detection. Summary of the invention

[0004] In view of this, the purpose of the embodiments of the present application is to provide a target detection model training method, device, electronic device and storage medium, which can improve the problem of low accuracy and poor effect of traditional target recognition models in detecting small targets in video satellite images.

[0005] In order to achieve the above technical objectives, the technical solutions adopted in this application are as follows:

[0006] In a first aspect, an embodiment of the present application provides a method for training a target detection model, the method comprising:

[0007] Acquire a data set for model training, wherein the data set includes a plurality of sample images and the type and coordinates of a real target corresponding to each sample image;

[0008] Inputting the data set into a feature extraction network in an initial detection model to be trained, and for each sample image, performing multi-scale feature extraction and attention feature extraction on the sample image through the feature extraction network to obtain a plurality of multi-scale feature maps corresponding to each sample image;

[0009] For each of the sample images, input each of the multi-scale feature maps into a detection head in the initial detection model to obtain a detection result corresponding to the sample image, wherein the detection result includes the type and coordinates of the target to be detected in the sample image;

[0010] Calculating the loss value between each of the to-be-detected targets and the category and coordinates of each of the to-be-detected targets and the type and coordinates of the corresponding real target, wherein the loss value represents the error between the to-be-detected target and the real target;

[0011] Determine whether the loss value meets the preset model convergence condition. If the loss value meets the preset model convergence condition, the initial detection model that meets the preset model convergence condition is used as the target detection model. If the loss value does not meet the preset model convergence condition, adjust the parameters of the initial detection model and train the initial detection model according to the adjusted parameters until the loss value meets the preset model convergence condition.

[0012] In combination with the first aspect, in some optional implementations, the data set is input into a feature extraction network in an initial detection model to be trained, and for each of the sample images, multi-scale feature extraction and attention feature extraction are performed on the sample image through the feature extraction network to obtain a plurality of multi-scale feature maps corresponding to each of the sample images, including:

[0013] Extracting feature maps of different scales in the sample image through the backbone network module;

[0014] Determine the feature map carrying the semantic features of the target to be detected in the feature maps of different scales as the target feature map, and input each of the target feature maps into the neck network module in the feature extraction network, so as to enhance the shallow features of the target feature map through the PAFPN structure in the neck network module, and obtain the enhanced target feature map corresponding to each of the target feature images;

[0015] For each of the enhanced target feature maps, the enhanced target feature maps are input into the trident spatial fusion attention submodule in the neck network module, so as to perform feature extraction on the enhanced target feature maps through the trident spatial fusion attention submodule to obtain a multi-scale feature map corresponding to each of the enhanced target feature maps.

[0016] In combination with the first aspect, in some optional implementations, for each of the sample images, the trident spatial attention submodule includes three spatial attention modules, one concatenate layer and one CBS module;

[0017] For each of the enhanced target feature maps, the enhanced target feature maps are input into the trident space fusion attention submodule in the neck network module, so as to perform feature extraction on the enhanced target feature maps through the trident space fusion attention submodule to obtain a multi-scale feature map corresponding to each of the enhanced target feature maps, including:

[0018] Through the three spatial attention modules, feature extraction is performed on different dimensions in the enhanced target feature map to obtain feature information corresponding to each dimension;

[0019] All feature information is concatenated as tensors through the concatenate layer and output through the CBS module to obtain the multi-scale feature map.

[0020] In combination with the first aspect, in some optional implementations, each of the spatial attention modules includes a first pooling layer, a second pooling layer, a first convolutional layer, and a first activation function;

[0021] Through the three spatial attention modules, feature extraction is performed on different dimensions in the enhanced target feature map to obtain feature information corresponding to each dimension, including:

[0022] Performing shape transformation on the enhanced target feature map to obtain three feature maps of different shapes;

[0023] Input the three feature maps of different shapes into the first pooling layer of each of the spatial attention modules to perform a maximum pooling operation, so as to obtain a first feature map corresponding to each of the feature maps;

[0024] Inputting the three feature maps of different shapes into the second pooling layer of each of the spatial attention modules respectively to perform an average pooling operation to obtain a second feature map corresponding to each of the feature maps;

[0025] For each of the spatial attention modules, performing tensor concatenation of the first feature map and the second feature map to obtain a third feature map;

[0026] Inputting the third feature map into the first convolution layer of each of the spatial attention modules for deep convolution, and calculating the weight coefficient corresponding to the third feature map after deep convolution by the first activation function;

[0027] The weight coefficient is multiplied by the enhanced target feature map to obtain the feature information.

[0028] In combination with the first aspect, in some optional implementations, calculating the loss value between the category and coordinates of each of the to-be-detected targets and the type and coordinates of the corresponding real target includes:

[0029] The loss value is determined by the loss function of the initial detection model, wherein the calculation formula of the loss value is as follows:

[0030]

[0031]

[0032]

[0033]

[0034] In the formula, CIoULoss represents the loss value, and IoU represents the difference between the predicted box B and the labeled true box B gt The intersection-over-union ratio, ρ 2 (b,b gt ) represents the center point b of the prediction box and the center point b of the true value box gt The Euclidean distance of the predicted box and the true box is , c represents the diagonal distance of the minimum enclosing area that can contain both the predicted box and the true value box, α represents the weight coefficient, v represents the similarity of the aspect ratio of the predicted box and the true value box, w represents the width of the predicted box, h represents the height of the predicted box, and w represents the height of the predicted box. gt Indicates the width of the truth box, h gt Indicates the height of the truth box.

[0035] In combination with the first aspect, in some optional implementations, the method further includes:

[0036] Obtain an image to be recognized;

[0037] The image to be identified is input into the target detection model to identify the target type of the target to be identified in the image to be identified.

[0038] In combination with the first aspect, in some optional implementations, before obtaining a data set for model training, the method further includes:

[0039] Get a sample image set;

[0040] Performing data enhancement processing on each initial sample image in the sample image set to obtain a sample image set after data enhancement;

[0041] The data-enhanced sample image set is divided into a training set and a test set according to a preset ratio to serve as the data set.

[0042] In a second aspect, an embodiment of the present application further provides a target detection model training device, the device comprising:

[0043] A first acquisition unit is used to acquire a data set for model training, wherein the data set includes a plurality of sample images and a type and coordinates of a real target corresponding to each sample image;

[0044] A feature extraction unit, used for inputting the data set into a feature extraction network in an initial detection model to be trained, and for each of the sample images, performing multi-scale feature extraction and attention feature extraction on the sample image through the feature extraction network to obtain a plurality of multi-scale feature maps corresponding to each of the sample images;

[0045] A detection unit, configured to input each of the multi-scale feature maps into a detection head in the initial detection model for each of the sample images, so as to obtain a detection result corresponding to the sample image, wherein the detection result includes a type and coordinates of a target to be detected in the sample image;

[0046] A second acquisition unit, used to calculate a loss value between the category and coordinates of each of the to-be-detected targets and the type and coordinates of the corresponding real target, wherein the loss value represents an error between the to-be-detected target and the real target;

[0047] A judgment unit is used to judge whether the loss value meets the preset model convergence condition. If the loss value meets the preset model convergence condition, the initial detection model that meets the preset model convergence condition is used as the target detection model. If the loss value does not meet the preset model convergence condition, the parameters of the initial detection model are adjusted, and the initial detection model is trained according to the adjusted parameters until the loss value meets the preset model convergence condition.

[0048] In a third aspect, an embodiment of the present application further provides an electronic device, comprising a processor and a memory coupled to each other, wherein a computer program is stored in the memory, and when the computer program is executed by the processor, the electronic device executes the above method.

[0049] In a fourth aspect, an embodiment of the present application further provides a computer-readable storage medium, wherein a computer program is stored in the computer-readable storage medium. When the computer program is executed on a computer, the computer executes the above method.

[0050] The invention adopting the above technical solution has the following advantages:

[0051] In the technical solution provided by the present application, a data set for model training is first obtained, and a sample image in the data set is input into a feature extraction network in the initial detection model to be trained, so as to perform multi-scale feature extraction and attention feature extraction on the sample image through the feature extraction network to obtain a multi-scale feature map, and then the multi-scale feature map is input into the detection head in the initial detection model to obtain the detection result corresponding to the sample image, and then the loss value between the target to be detected and the real target in the detection result is obtained. When the loss value meets the preset model convergence condition, the initial detection model is determined to be a target detection model. In this way, the problem of low accuracy and poor effect of traditional target recognition models in detecting small targets in video satellite images can be improved. BRIEF DESCRIPTION OF THE DRAWINGS

[0052] The present application may be further described by the non-limiting embodiments given in the accompanying drawings. It should be understood that the following drawings only illustrate certain embodiments of the present application and therefore should not be regarded as limiting the scope. For those of ordinary skill in the art, other relevant drawings may be obtained based on these drawings without creative effort.

[0053] Figure 1 A block diagram of an electronic device provided in an embodiment of the present application.

[0054] Figure 2 A flowchart of a target detection model training method provided in an embodiment of the present application.

[0055] Figure 3 This is one of the structural schematic diagrams of the initial detection model provided in the embodiment of the present application.

[0056] Figure 4 The second structural diagram of the initial detection model provided in the embodiment of the present application.

[0057] Figure 5 A schematic diagram of the structure of the Trident spatial fusion attention submodule provided in an embodiment of the present application.

[0058] Figure 6 A block diagram of a target detection model training device provided in an embodiment of the present application.

[0059] Icon: 100 - electronic device; 101 - processor; 102 - memory; 200 - target detection model training device; 210 - first acquisition unit; 220 - feature extraction unit; 230 - detection unit; 240 - second acquisition unit; 250 - judgment unit. DETAILED DESCRIPTION

[0060] The present application will be described in detail below in conjunction with the accompanying drawings and specific embodiments. It should be noted that in the drawings or descriptions, similar or identical parts use the same figure numbers, and the implementation methods not shown or described in the drawings are forms known to ordinary technicians in the relevant technical field. In the description of this application, the terms "first", "second", etc. are only used to distinguish the description and cannot be understood as indicating or implying relative importance.

[0061] Please refer to Figure 1 , an embodiment of the present application provides an electronic device 100 that may include a processor 101 and a memory 102. The memory 102 stores a computer program, and when the computer program is executed by the processor 101, the electronic device 100 can perform the corresponding steps in the following target detection model training method.

[0062] In this embodiment, the electronic device 100 can be a personal computer, a server, a laptop computer, etc., which is used to obtain a data set for model training, and input the sample images in the data set into the initial detection model to be trained to obtain a multi-scale feature map, and then input the multi-scale feature map into the detection head of the initial detection model to obtain the detection result, and adjust the parameters of the initial detection model according to the loss value between the target to be detected and the real target in the detection result and the preset model convergence condition.

[0063] Please refer to Figure 2 The present application also provides a target detection model training method, which can be applied to the above electronic device 100, and each step in the method is executed or implemented by the electronic device 100. The target detection model training method can include the following steps:

[0064] Step 110, obtaining a data set for model training, wherein the data set includes a plurality of sample images and the type and coordinates of a real target corresponding to each sample image;

[0065] Step 120, inputting the data set into a feature extraction network in the initial detection model to be trained, and for each of the sample images, performing multi-scale feature extraction and attention feature extraction on the sample image through the feature extraction network to obtain a plurality of multi-scale feature maps corresponding to each of the sample images;

[0066] Step 130: for each of the sample images, input each of the multi-scale feature maps into a detection head in the initial detection model to obtain a detection result corresponding to the sample image, wherein the detection result includes the type and coordinates of the target to be detected in the sample image;

[0067] Step 140, calculating a loss value between the category and coordinates of each of the to-be-detected targets and the category and coordinates of the corresponding real target, wherein the loss value represents an error between the to-be-detected target and the real target;

[0068] Step 150, determine whether the loss value meets the preset model convergence condition. If the loss value meets the preset model convergence condition, the initial detection model that meets the preset model convergence condition is used as the target detection model. If the loss value does not meet the preset model convergence condition, adjust the parameters of the initial detection model and train the initial detection model according to the adjusted parameters until the loss value meets the preset model convergence condition.

[0069] In the technical solution provided by the present application, a data set for model training is first obtained, and a sample image in the data set is input into a feature extraction network in the initial detection model to be trained, so as to perform multi-scale feature extraction and attention feature extraction on the sample image through the feature extraction network to obtain a multi-scale feature map, and then the multi-scale feature map is input into the detection head in the initial detection model to obtain the detection result corresponding to the sample image, and then the loss value between the target to be detected and the real target in the detection result is obtained. When the loss value meets the preset model convergence condition, the initial detection model is determined to be a target detection model. In this way, the problem of low accuracy and poor effect of traditional target recognition models in detecting small targets in video satellite images can be improved.

[0070] The following is a detailed description of the various steps of the target detection model training method, as follows:

[0071] Prior to step 110, the method may include:

[0072] Get a sample image set;

[0073] Performing data enhancement processing on each initial sample image in the sample image set to obtain a sample image set after data enhancement;

[0074] The data-enhanced sample image set is divided into a training set and a test set according to a preset ratio to serve as the data set.

[0075] In this embodiment, the sample image set may be a plurality of single-frame images captured from a video stream generated by satellite remote sensing of the ground, or may be a plurality of camera images obtained by shooting satellite monitoring screens. The sample image set may be acquired by pre-storing the sample image set in the memory 102 of the electronic device 100, and then calling the sample image set from the memory 102 based on the user's operation instruction, or by capturing or shooting images in real time based on the user's operation instruction and generating the sample image set. The acquisition method of the sample image set is not specifically limited here.

[0076] Understandably, model training usually requires a large number of samples to support the parameter tuning of the model during the training process, and the initial neural network in the model will also recognize slightly changed image samples as unique and different images. Therefore, by making slight changes to the existing data set, such as rotation, flipping, scaling, cropping, shifting and other conventional methods, the initial sample images obtained are enhanced to increase the amount of sample data.

[0077] In this embodiment, in addition to using the above conventional method to perform data enhancement on the initial sample image, Mosaic data enhancement is also used. This method randomly selects 4 pictures, randomly crops and scales them, and then randomly arranges and splices them to form an image. While enriching the data set, it increases the small sample target and improves the training speed of the network. When performing the normalization operation, the data of 4 pictures will be calculated at one time to reduce the memory requirements of the model.

[0078] In this embodiment, the preset ratio can be flexibly set according to user needs, such as 8:2, 7:3, etc. Exemplarily, the sample image set after data enhancement can be divided into a training set and a test set at a ratio of 8:2, and the sum of the training set and the test set is used as the data set to train the initial detection model.

[0079] In step 110, a pre-stored data set may be called based on a user instruction, or an initial sample image may be acquired in real time and a data set may be generated through the aforementioned steps to serve as a data source for subsequent initial detection model training.

[0080] In step 120, the data set is input into a feature extraction network in an initial detection model to be trained, and for each sample image, multi-scale feature extraction and attention feature extraction are performed on the sample image through the feature extraction network to obtain a plurality of multi-scale feature maps corresponding to each sample image, which may include:

[0081] Extracting feature maps of different scales in the sample image through the backbone network module;

[0082] Determine the feature map carrying the semantic features of the target to be detected in the feature maps of different scales as the target feature map, and input each of the target feature maps into the neck network module in the feature extraction network, so as to enhance the shallow features of the target feature map through the PAFPN structure in the neck network module, and obtain the enhanced target feature map corresponding to each of the target feature images;

[0083] For each of the enhanced target feature maps, the enhanced target feature maps are input into the trident spatial fusion attention submodule in the neck network module, so as to perform feature extraction on the enhanced target feature maps through the trident spatial fusion attention submodule to obtain a multi-scale feature map corresponding to each of the enhanced target feature maps.

[0084] In this embodiment, refer to Figure 3 , feature maps of different scales are generated through the backbone network, and P3, P4 and P5 carrying semantic features are selected from the feature maps of different scales as target feature maps and input into the neck network module. Then, the PAFPN structure in the neck network module enhances the shallow features of the target feature maps to obtain enhanced target feature maps N3, N4 and N5. Then, the enhanced target feature maps are input into the Trident spatial fusion attention submodule, and attention features are extracted from the enhanced target feature maps to obtain multi-scale feature maps C3, C4 and C5 as the input of the subsequent detection head.

[0085] In this embodiment, the backbone network adopts the CSPDarkNet structure. CSPDarkNet solves the problem of gradient information duplication in model training, integrates the gradient changes from beginning to end into the feature map, reduces the number of model parameters and FLOPS (floating point operations, which can be understood as the amount of calculation), and reduces the model size while ensuring the inference speed and accuracy.

[0086] In this embodiment, refer to Figure 4 ,In the figure, the CBS module includes the second convolutional layer, the BN layer, and the second activation function, represents two downsampling operations, Represents tensor concatenation in the channel dimension, C3 includes 3 CBS modules plus a CSPbottleneck module, and TSFA represents the Trident Spatial Fusion Attention submodule. Feature maps P1, P2, P3, P4, and P5 of different scales of the sample image are extracted through the backbone network. Although P1 and P2 have high resolution and rich physical information (position information), adding P1 or P2 will generate huge computational complexity, and P1 and P2 lack semantic features. Using them as the input of the detection head is prone to wrong predictions, which reduces the detection performance of the model. Therefore, P3, P4, and P5 carrying semantic features are used as target feature maps as inputs to the subsequent neck network module and detection head.

[0087] The target feature map is input into the neck network module, and the PAFPN structure in the neck network module enhances the shallow features of the target feature map to obtain enhanced target feature maps N3, N4, and N5. The PAFPN structure introduces a bottom-up path based on the feature pyramid network (FPN). Starting from the bottom layer, it gradually merges the new features obtained in the current layer and the previous layer. Taking P3 as an example, P3 and P4 that have been downsampled twice are tensor spliced ​​in the channel dimension, and the output is N3 after passing through 3 C3 modules, and so on, until the highest layer, the enhanced feature maps N4 and N5 are output in turn, so that the shallow features can be better transmitted upward in the feature pyramid, thereby enhancing the expressive power of the feature pyramid.

[0088] After obtaining the enhanced target feature map, the enhanced target feature maps N3, N4, and N5 are respectively input into the Trident spatial fusion attention submodule (i.e., the TSFA module shown in the figure), and the enhanced target feature map is subjected to feature extraction through the Trident spatial fusion attention submodule to obtain the multi-scale feature maps C3, C4, and C5.

[0089] In this embodiment, refer to Figure 5 For each sample image, the Trident spatial attention submodule includes three spatial attention modules, one concatenate layer (i.e., ) and 1 CBS module, each of the spatial attention modules includes a first pooling layer (i.e., MaxPool shown in the figure), a second pooling layer (i.e., AvgPool shown in the figure), a first convolutional layer (i.e., DWConv shown in the figure), and a first activation function (i.e., ), the CBS module includes the second convolutional layer (i.e., Conv shown in the figure), the BN layer, and the second activation function (i.e., SiLU shown in the figure).

[0090] For each of the enhanced target feature maps, inputting the enhanced target feature map into the trident space fusion attention submodule in the neck network module, so as to perform feature extraction on the enhanced target feature map through the trident space fusion attention submodule to obtain a multi-scale feature map corresponding to each of the enhanced target feature maps, may include:

[0091] Through the three spatial attention modules, feature extraction is performed on different dimensions in the enhanced target feature map to obtain feature information corresponding to each dimension;

[0092] All feature information is concatenated as tensors through the concatenate layer and output through the CBS module to obtain the multi-scale feature map.

[0093] Among them, by using the three spatial attention modules, feature extraction is performed on different dimensions in the enhanced target feature map to obtain feature information corresponding to each dimension, which may include:

[0094] Performing shape transformation on the enhanced target feature map to obtain three feature maps of different shapes;

[0095] Input the three feature maps of different shapes into the first pooling layer of each of the spatial attention modules to perform a maximum pooling operation, so as to obtain a first feature map corresponding to each of the feature maps;

[0096] Inputting the three feature maps of different shapes into the second pooling layer of each of the spatial attention modules respectively to perform an average pooling operation to obtain a second feature map corresponding to each of the feature maps;

[0097] For each of the spatial attention modules, performing tensor concatenation of the first feature map and the second feature map to obtain a third feature map;

[0098] Inputting the third feature map into the first convolution layer of each of the spatial attention modules for deep convolution, and calculating the weight coefficient corresponding to the third feature map after deep convolution by the first activation function;

[0099] The weight coefficient is multiplied by the enhanced target feature map to obtain the feature information.

[0100] In this embodiment, again referring to Figure 5 , any enhanced target feature map is recorded as feature map F. When feature map F is used as the input of the Trident spatial fusion attention submodule, it includes three dimensions (C, H, W), namely channel dimension, high latitude and wide latitude. Spatial attention mechanisms (i.e., spatial attention modules) are added to the feature maps of the corresponding shapes in the three dimensions. Taking the C channel dimension as an example, the feature map F with a shape of C×H×W is subjected to maximum pooling and average pooling to obtain the texture information and feature information of the feature map F. Among them, the maximum pooling is to take the maximum value of the pixels in the pooling area. The first feature map obtained in this way is more sensitive to texture feature information; the average pooling is to take the average value of the image in the pooling area. The second feature map obtained in this way is more sensitive to background information. Then the results after the two poolings (i.e., the first feature map and the second feature map) are tensor spliced ​​to obtain the attention output F of the C channel dimension. C (i.e. the third characteristic graph), as shown below:

[0101] F C =Concat(AcgPool(F C×H×W ),MaxPool(F C×H×W ))

[0102] In the formula, Concat() represents tensor concatenation, AvgPool() represents average pooling, MaxPool() represents maximum pooling, and C×H×W represents the shape of the feature map F. The two 1×H×W images obtained by average pooling and maximum pooling are concatenated into a 2×H×W feature map to obtain the attention output F of the C channel dimension. C .

[0103] Similarly, the shape of the feature map F is transformed into W×H×C and H×C×W, and average pooling, maximum pooling and tensor concatenation are performed to obtain the attention output F with a width of W. W and H high-dimensional attention output F H ,as follows:

[0104] F W =Concat(AvgPool(F W×H×C ),MaxPool(F W×H×C ))

[0105] F H =Concat(AvgPool(F H×C×W ),MaxPool(F H×C×W ))

[0106] Then, the attention output of each dimension (i.e., F C 、F W and F H ) performs a deep convolution operation (DWConv, Depth-wise Conv). Compared with conventional convolution operations, deep convolution can reduce the amount of convolution calculations. Then, after a Sigmoid activation function (i.e., the first activation function), the weight coefficient corresponding to each dimension after the deep convolution is calculated by the first activation function, and then the weight coefficient is multiplied by the feature map F to obtain the scaled attention feature information F′ of different dimensions. C , F′ H and F′ W ,as follows:

[0107] F′ C =F×sigmoid(DWConv(F C ))

[0108] F′ W =F×sigmoid(DWConv(F W ))

[0109] F′ H =F×sigmoid(DWConv(F H ))

[0110] After obtaining the attention feature information, all the attention feature information is concatenated through the above concatenate layer and output through the CBS module to obtain a multi-scale feature map, as follows:

[0111] F refined =CBS(Concat(F′ W ,F′ C ,F′ H ))

[0112] In the formula, F refined Represents the feature map F after spatial attention refinement, that is, the multi-scale feature map.

[0113] In step 130, the detection head of the initial detection model is used to divide the above multi-scale feature map into grids, each grid is responsible for detecting a set of predefined anchor boxes, each anchor box carries the corresponding boundary coordinates and size, and the confidence score of the target to be detected within the anchor box. Then, a prediction box of each target to be detected is generated based on the predefined anchor box, and the prediction box carries the type and coordinates of the target to be detected, that is, the detection result.

[0114] In this embodiment, after obtaining the multi-scale feature map, the multi-scale feature map is input into the detection head of the initial detection model, and the detection head detects and obtains the detection result corresponding to the sample image. The detection result can be the type and position coordinates of small targets in the satellite image, such as cars, airplanes, ships, gasoline barrels, etc.

[0115] In this embodiment, the structure of the detection head can adopt a detection head of any network model that can realize target detection, such as a detection head of a model structure such as YOLOv3, YOLOv5, etc.

[0116] In step 140, calculating the loss value between the category and coordinates of each of the to-be-detected targets and the type and coordinates of the corresponding real target may include:

[0117] The loss value is determined by the loss function of the initial detection model, wherein the calculation formula of the loss value is as follows:

[0118]

[0119]

[0120]

[0121]

[0122] In the formula, CIoULoss represents the loss value, and IoU represents the difference between the predicted box B and the labeled true box B gt The intersection-over-union ratio, ρ2 (b,b gt ) represents the center point b of the prediction box and the center point b of the true value box gt The Euclidean distance of the predicted box and the true box is , c represents the diagonal distance of the minimum enclosing area that can contain both the predicted box and the true value box, α represents the weight coefficient, v represents the similarity of the aspect ratio of the predicted box and the true value box, w represents the width of the predicted box, h represents the height of the predicted box, and w represents the height of the predicted box. gt Indicates the width of the truth box, h gt Indicates the height of the truth box.

[0123] In practical applications, in order to reduce the amount of data for model calculation and enhance the accuracy of model prediction, after the prediction box of the target to be detected is obtained, the detected detection box is screened by the DIoU-NMS algorithm to obtain the prediction box corresponding to the target to be detected, and then the loss value between the target to be detected and its corresponding real target is calculated by the above loss value calculation formula.

[0124] In step 150, the preset model convergence condition may refer to determining the loss value between the target to be detected and the true target after each iterative training through the loss function during the model training process. When the loss value fluctuates within a preset number of iterations and tends to be stable (that is, the extreme difference is within a preset range), it is determined that the loss function converges.

[0125] In this embodiment, the loss value between each target to be detected and its corresponding real target after each iterative training is calculated through step S140, and then it is determined whether the numerical fluctuation of the loss value tends to be stable within a preset number of iterations. If so, it is determined that the initial detection model meets the preset model convergence condition (that is, the loss function of the initial detection model converges), and the initial detection model is determined to be the target detection model; if not, the model parameters are continuously adjusted according to the loss value for iterative training until the loss value meets the preset model convergence condition.

[0126] Exemplarily, the preset interval of the number of iterations can be 500 times, and the preset range of the range of the loss value can be 0.06. After the model is iteratively trained 500 times, if in the first 500 training times, the range of the 500 loss values ​​corresponding to the model is still greater than 0.06, then the model parameters are continued to be adjusted according to the loss function for iterative training. If after the model is iteratively trained 1000 times, the range of the loss values ​​corresponding to the model in 501-1000 training times is less than or equal to 0.06, then it is determined that the initial detection model meets the preset model convergence conditions, and the initial detection model is determined to be the target detection model.

[0127] As an optional implementation, the method may further include:

[0128] Obtain an image to be recognized;

[0129] The image to be identified is input into the target detection model to identify the target type of the target to be identified in the image to be identified.

[0130] In this embodiment, after the training of the initial detection model is completed, the trained initial detection model is used as the target detection model. A single frame image in the video stream generated by satellite remote sensing ground is captured as the image to be identified, and then the image to be identified is input into the target detection model. The target detection model is used to identify the type and coordinate position of small targets in the image to be identified, such as airplanes in airports, ships in ports, vehicles on roads, etc.

[0131] Please refer to Figure 6 The present application also provides a target detection model training device 200, which includes at least one software function module that can be stored in the memory 102 in the form of software or firmware or fixed in the operating system (OS) of the electronic device 100. The processor 101 is used to execute the executable modules stored in the memory 102, such as the software function modules and computer programs included in the target detection model training device 200.

[0132] The target detection model training device 200 includes a first acquisition unit 210, a feature extraction unit 220, a detection unit 230, a second acquisition unit 240 and a judgment unit 250. The functions of each unit may be as follows:

[0133] A first acquisition unit 210 is used to acquire a data set for model training, wherein the data set includes a plurality of sample images and a type and coordinates of a real target corresponding to each sample image;

[0134] A feature extraction unit 220 is used to input the data set into a feature extraction network in the initial detection model to be trained, and for each of the sample images, perform multi-scale feature extraction and attention feature extraction on the sample image through the feature extraction network to obtain a plurality of multi-scale feature maps corresponding to each of the sample images;

[0135] A detection unit 230, configured to input each of the multi-scale feature maps into a detection head in the initial detection model for each of the sample images, so as to obtain a detection result corresponding to the sample image, wherein the detection result includes a type and coordinates of a target to be detected in the sample image;

[0136] A second acquisition unit 240 is used to calculate a loss value between the category and coordinates of each of the to-be-detected targets and the type and coordinates of the corresponding real target, wherein the loss value represents an error between the to-be-detected target and the real target;

[0137] The judgment unit 250 is used to judge whether the loss value meets the preset model convergence condition. If the loss value meets the preset model convergence condition, the initial detection model that meets the preset model convergence condition is used as the target detection model. If the loss value does not meet the preset model convergence condition, the parameters of the initial detection model are adjusted, and the initial detection model is trained according to the adjusted parameters until the loss value meets the preset model convergence condition.

[0138] Optionally, the feature extraction unit 220 is further configured to:

[0139] Extracting feature maps of different scales in the sample image through the backbone network module;

[0140] Determine the feature map carrying the semantic features of the target to be detected in the feature maps of different scales as the target feature map, and input each of the target feature maps into the neck network module in the feature extraction network, so as to enhance the shallow features of the target feature map through the PAFPN structure in the neck network module, and obtain the enhanced target feature map corresponding to each of the target feature images;

[0141] For each of the enhanced target feature maps, the enhanced target feature maps are input into the trident spatial fusion attention submodule in the neck network module, so as to perform feature extraction on the enhanced target feature maps through the trident spatial fusion attention submodule to obtain a multi-scale feature map corresponding to each of the enhanced target feature maps.

[0142] Optionally, for each of the sample images, the Trident spatial attention submodule includes 3 spatial attention modules, 1 concatenate layer and 1 CBS module;

[0143] The feature extraction unit 220 is also used for:

[0144] Through the three spatial attention modules, feature extraction is performed on different dimensions in the enhanced target feature map to obtain feature information corresponding to each dimension;

[0145] All feature information is concatenated as tensors through the concatenate layer and output through the CBS module to obtain the multi-scale feature map.

[0146] Optionally, each of the spatial attention modules includes a first pooling layer, a second pooling layer, a first convolutional layer and a first activation function;

[0147] The feature extraction unit 220 is also used for:

[0148] Performing shape transformation on the enhanced target feature map to obtain three feature maps of different shapes;

[0149] Inputting the three feature maps of different shapes into the first pooling layer of each of the spatial attention modules respectively to perform a maximum pooling operation to obtain a first feature map corresponding to each of the feature maps;

[0150] Inputting the three feature maps of different shapes into the second pooling layer of each of the spatial attention modules respectively to perform an average pooling operation to obtain a second feature map corresponding to each of the feature maps;

[0151] For each of the spatial attention modules, performing tensor concatenation of the first feature map and the second feature map to obtain a third feature map;

[0152] Inputting the third feature map into the first convolution layer of each of the spatial attention modules for deep convolution, and calculating the weight coefficient corresponding to the third feature map after deep convolution by the first activation function;

[0153] The weight coefficient is multiplied by the enhanced target feature map to obtain the feature information.

[0154] Optionally, the second acquiring unit 240 is further configured to:

[0155] The loss value is determined by the loss function of the initial detection model, wherein the calculation formula of the loss value is as follows:

[0156]

[0157]

[0158]

[0159]

[0160] In the formula, CIoULoss represents the loss value, and IoU represents the difference between the predicted box B and the labeled true box B gt The intersection-over-union ratio, ρ 2 (b,b gt ) represents the center point b of the prediction box and the center point b of the true value box gt The Euclidean distance of the predicted box and the true box is , c represents the diagonal distance of the minimum enclosing area that can contain both the predicted box and the true value box, α represents the weight coefficient, v represents the similarity of the aspect ratio of the predicted box and the true value box, w represents the width of the predicted box, h represents the height of the predicted box, and w represents the height of the predicted box. gt Indicates the width of the truth box, h gt Indicates the height of the truth box.

[0161] Optionally, the target detection model training device 200 further includes:

[0162] A third acquisition module is used to acquire an image to be identified;

[0163] The recognition module is used to input the image to be recognized into the target detection model to recognize the target type of the target to be recognized in the image to be recognized.

[0164] In this embodiment, the processor 101 may be an integrated circuit chip having signal processing capabilities. The processor 101 may be a general-purpose processor. For example, the processor 101 may be a central processing unit (CPU), a digital signal processor (DSP), an application specific integrated circuit (ASIC), a field programmable gate array (FPGA) or other programmable logic devices, discrete gates or transistor logic devices, discrete hardware components, and may implement or execute the methods, steps and logic block diagrams disclosed in the embodiments of the present application.

[0165] The memory 102 may be, but is not limited to, a random access memory, a read-only memory, a programmable read-only memory, an erasable programmable read-only memory, an electrically erasable programmable read-only memory, etc. In this embodiment, the memory 102 may be used to store a data set, an initial detection model, a loss function, a loss value, a preset model convergence condition, etc. Of course, the memory 102 may also be used to store a program, and the processor 101 executes the program after receiving an execution instruction.

[0166] Understandably, Figure 1 The structure of the electronic device 100 shown in FIG. 1 is only a schematic diagram of a structure. The electronic device 100 may also include Figure 1 More components are shown. Figure 1 Each component shown in the figure can be implemented by hardware, software or a combination thereof.

[0167] It should be noted that those skilled in the art can clearly understand that, for the convenience and brevity of description, the specific working process of the electronic device 100 described above can refer to the corresponding process of each step in the aforementioned method, and will not be elaborated herein.

[0168] The embodiment of the present application also provides a computer-readable storage medium. The computer-readable storage medium stores a computer program, and when the computer program is run on a computer, the computer executes the target detection model training method described in the above embodiment.

[0169] Through the description of the above implementation methods, technical personnel in this field can clearly understand that the present application can be implemented by hardware, and can also be implemented by means of software plus a necessary general hardware platform. Based on such an understanding, the technical solution of the present application can be embodied in the form of a software product, which can be stored in a non-volatile storage medium (which can be a CD-ROM, a USB flash drive, a mobile hard disk, etc.), including a number of instructions for enabling a computer device (which can be a personal computer, a server, or a network device, etc.) to execute the methods described in each implementation scenario of the present application.

[0170] In summary, the embodiments of the present application provide a method, device, electronic device and storage medium for training a target detection model. In the present technical solution, a data set for model training is first obtained, and a sample image in the data set is input into a feature extraction network in an initial detection model to be trained, so as to perform multi-scale feature extraction and attention feature extraction on the sample image through the feature extraction network to obtain a multi-scale feature map, and then the multi-scale feature map is input into the detection head in the initial detection model to obtain a detection result corresponding to the sample image, and then the loss value between the target to be detected and the real target in the detection result is obtained. When the loss value meets the preset model convergence condition, the initial detection model is determined to be a target detection model. In this way, the problem of low accuracy and poor effect of traditional target recognition models in detecting small targets in video satellite images can be improved.

[0171] In the embodiments provided in the present application, it should be understood that the disclosed device, system and method can also be implemented in other ways. The above-described device, system and method embodiments are merely schematic, for example, the flowchart and block diagram in the accompanying drawings show the possible architecture, function and operation of the system, method and computer program product according to the multiple embodiments of the present application. In this regard, each box in the flowchart or block diagram can represent a part of a module, a program segment or a code, and a part of the module, program segment or code contains one or more executable instructions for implementing the specified logical function. It should also be noted that each box in the block diagram and / or the flowchart, and the combination of the boxes in the block diagram and / or the flowchart can be implemented by a dedicated hardware-based system that performs the specified function or action, or can be implemented by a combination of dedicated hardware and computer instructions. In addition, each functional module in each embodiment of the present application can be integrated together to form an independent part, or each module can exist separately, or two or more modules can be integrated to form an independent part.

[0172] The above description is only an embodiment of the present application and is not intended to limit the protection scope of the present application. For those skilled in the art, the present application may have various modifications and variations. Any modification, equivalent replacement, improvement, etc. made within the spirit and principle of the present application shall be included in the protection scope of the present application.

Claims

1. A target detection model training method, characterized in that: The method comprises: Acquire a data set for model training, wherein the data set includes a plurality of sample images and the type and coordinates of a real target corresponding to each sample image; Inputting the data set into a feature extraction network in an initial detection model to be trained, and for each sample image, performing multi-scale feature extraction and attention feature extraction on the sample image through the feature extraction network to obtain a plurality of multi-scale feature maps corresponding to each sample image; For each of the sample images, input each of the multi-scale feature maps into a detection head in the initial detection model to obtain a detection result corresponding to the sample image, wherein the detection result includes the type and coordinates of the target to be detected in the sample image; Calculating a loss value between the category and coordinates of each of the to-be-detected targets and the corresponding category and coordinates of the real target, wherein the loss value represents an error between the category and coordinates of the to-be-detected targets and the category and coordinates of the real target; Determine whether the loss value meets the preset model convergence condition. If the loss value meets the preset model convergence condition, the initial detection model that meets the preset model convergence condition is used as the target detection model. If the loss value does not meet the preset model convergence condition, adjust the parameters of the initial detection model and train the initial detection model according to the adjusted parameters until the loss value meets the preset model convergence condition.

2. The method according to claim 1, characterized in that The data set is input into a feature extraction network in an initial detection model to be trained. For each sample image, multi-scale feature extraction and attention feature extraction are performed on the sample image through the feature extraction network to obtain a plurality of multi-scale feature maps corresponding to each sample image, including: Extracting feature maps of different scales in the sample image through the backbone network module; Determine the feature map carrying the semantic features of the target to be detected in the feature maps of different scales as the target feature map, and input each of the target feature maps into the neck network module in the feature extraction network, so as to enhance the shallow features of the target feature map through the PAFPN structure in the neck network module, and obtain the enhanced target feature map corresponding to each of the target feature images; For each of the enhanced target feature maps, the enhanced target feature maps are input into the trident spatial fusion attention submodule in the neck network module, so as to perform feature extraction on the enhanced target feature maps through the trident spatial fusion attention submodule to obtain a multi-scale feature map corresponding to each of the enhanced target feature maps.

3. The method according to claim 2, characterized in that For each of the sample images, the Trident spatial attention submodule includes three spatial attention modules, one concatenate layer and one CBS module; For each of the enhanced target feature maps, the enhanced target feature maps are input into the trident space fusion attention submodule in the neck network module, so as to perform feature extraction on the enhanced target feature maps through the trident space fusion attention submodule to obtain a multi-scale feature map corresponding to each of the enhanced target feature maps, including: Through the three spatial attention modules, feature extraction is performed on different dimensions in the enhanced target feature map to obtain feature information corresponding to each dimension; All feature information is concatenated as tensors through the concatenate layer and output through the CBS module to obtain the multi-scale feature map.

4. The method according to claim 3, characterized in that Each of the spatial attention modules comprises a first pooling layer, a second pooling layer, a first convolutional layer and a first activation function; Through the three spatial attention modules, feature extraction is performed on different dimensions in the enhanced target feature map to obtain feature information corresponding to each dimension, including: Performing shape transformation on the enhanced target feature map to obtain three feature maps of different shapes; Input the three feature maps of different shapes into the first pooling layer of each of the spatial attention modules to perform a maximum pooling operation, so as to obtain a first feature map corresponding to each of the feature maps; Input the three feature maps of different shapes into the second pooling layer of each of the spatial attention modules respectively to perform an average pooling operation to obtain a second feature map corresponding to each of the feature maps; For each of the spatial attention modules, performing tensor concatenation of the first feature map and the second feature map to obtain a third feature map; Inputting the third feature map into the first convolution layer of each of the spatial attention modules for deep convolution, and calculating the weight coefficient corresponding to the third feature map after deep convolution by the first activation function; The weight coefficient is multiplied by the enhanced target feature map to obtain the feature information.

5. The method according to claim 1, characterized in that Calculating the loss value between the category and coordinates of each of the to-be-detected targets and the type and coordinates of the corresponding real target, including: The loss value is determined by the loss function of the initial detection model, wherein the calculation formula of the loss value is as follows: In the formula, CIoU Loss represents the loss value, and IoU represents the difference between the predicted box B and the labeled true box B gt The intersection-over-union ratio, ρ 2 (b,b gt ) represents the center point b of the prediction box and the center point b of the true value box gt The Euclidean distance of the predicted box and the true box is , c represents the diagonal distance of the minimum enclosing area that can contain both the predicted box and the true value box, α represents the weight coefficient, v represents the similarity of the aspect ratio of the predicted box and the true value box, w represents the width of the predicted box, h represents the height of the predicted box, and w represents the height of the predicted box. gt Indicates the width of the truth box, h gt Indicates the height of the truth box.

6. The method according to claim 1, characterized in that The method further comprises: Obtain an image to be recognized; The image to be identified is input into the target detection model to identify the target type of the target to be identified in the image to be identified.

7. The method according to claim 1, characterized in that Before obtaining a data set for model training, the method further includes: Get a sample image set; Performing data enhancement processing on each initial sample image in the sample image set to obtain a data enhanced sample image set; The data-enhanced sample image set is divided into a training set and a test set according to a preset ratio to serve as the data set.

8. A target detection model training device, characterized in that: The device comprises: A first acquisition unit is used to acquire a data set for model training, wherein the data set includes a plurality of sample images and a type and coordinates of a real target corresponding to each sample image; A feature extraction unit, used for inputting the data set into a feature extraction network in an initial detection model to be trained, and for each of the sample images, performing multi-scale feature extraction and attention feature extraction on the sample image through the feature extraction network to obtain a plurality of multi-scale feature maps corresponding to each of the sample images; A detection unit, configured to input each of the multi-scale feature maps into a detection head in the initial detection model for each of the sample images, so as to obtain a detection result corresponding to the sample image, wherein the detection result includes a type and coordinates of a target to be detected in the sample image; A second acquisition unit, used to calculate a loss value between the category and coordinates of each of the to-be-detected targets and the type and coordinates of the corresponding real target, wherein the loss value represents an error between the to-be-detected target and the real target; A judgment unit is used to judge whether the loss value meets the preset model convergence condition. If the loss value meets the preset model convergence condition, the initial detection model that meets the preset model convergence condition is used as the target detection model. If the loss value does not meet the preset model convergence condition, the parameters of the initial detection model are adjusted, and the initial detection model is trained according to the adjusted parameters until the loss value meets the preset model convergence condition.

9. An electronic device, characterized in that: The electronic device comprises a processor and a memory coupled to each other, wherein the memory stores a computer program. When the computer program is executed by the processor, the electronic device executes the method as claimed in any one of claims 1 to 7.

10. A computer-readable storage medium, characterized in that: The computer-readable storage medium stores a computer program, and when the computer program is executed on a computer, the computer is enabled to execute the method according to any one of claims 1 to 7.

Images