Remote sensing small target detection method based on YOLOv5

Abstract

The invention relates to a remote sensing small target detection method based on YOLOv5, and belongs to the field of target detection, and the method comprises the steps: S1, collecting and preprocessing a data set, and dividing the data set into a training set, a verification set and a test set; s2, converting a label file in the data set into a YOLOv5 format; s3, creating a configuration file of the data set; s4, constructing an RStargetNet network: on the basis of the YOLOv5 network, using VanilaNet as a backbone network, adding a SENet channel attention mechanism, adopting BCEloss for confidence loss and category loss, and adopting CIoUloss for positioning loss; s5, the RStargetNet network is trained, verified and tested, and an optimal model is obtained; and S6, predicting the picture input into the network by using the optimal model, obtaining the category name, the confidence score and the target frame coordinate information of each detection target, and displaying the category name, the confidence score and the target frame coordinate information in the picture.

Description

Technical Field

[0001] The present invention belongs to the field of target detection technology and relates to a remote sensing small target detection method based on YOLOv5 Background Art

[0002] Object detection has been a long-standing problem in remote sensing images and computer vision. It is usually defined as identifying the location of target objects in an input image as well as identifying the object category. Automatic object detection has been widely used in many practical applications such as hazard detection, environmental monitoring, change detection, urban planning, etc. In the past few decades, object detection has been extensively studied and a large number of methods have been developed to detect artificial objects (such as vehicles, buildings, roads, bridges, etc.) and natural objects (such as lakes, coasts, forests, etc.) in remote sensing images. Machine learning-based methods have strong robustness in feature extraction and object classification and have been widely studied by many recent methods to achieve significant progress in this problem. In the past few years, few-shot learning has been widely studied in the field of computer vision for the tasks of scene classification, image segmentation, and object detection. In remote sensing images, the sizes of objects can be very different, and the spatial resolution of remote sensing images can also be very different, which makes this problem even more challenging when only a few annotated samples are provided.

[0003] With the advancement of remote sensing image technology, the feature information of remote sensing images is becoming more and more. The remote sensing image detection method based on deep learning obtains abstract semantic information in the image through training on the algorithm model through a large number of parameters, so that target recognition has higher detection accuracy and better adaptability. Summary of the invention

[0004] In view of this, the purpose of the present invention is to provide a remote sensing small target detection method based on YOLOv5, using VanillaNet to optimize the network structure of YOLOv5, adding SENet channel attention mechanism, and gradually eliminating nonlinear layers during training so that they can be easily merged while saving inference data. This allows the original network to recognize small targets in remote sensing images more efficiently.

[0005] In order to achieve the above object, the present invention provides the following technical solutions:

[0006] A remote sensing small target detection method based on YOLOv5, comprising the following steps:

[0007] S1: Collect remote sensing image data sets, preprocess the data sets, and divide them into training sets, validation sets, and test sets;

[0008] S2: Convert the label file in the dataset to YOLOv5 format;

[0009] S3: Create a configuration file for the dataset;

[0010] S4: Build RStargetNet network: Based on the YOLOv5 network, use VanillaNet to replace the CSPDarknet53 backbone network, add SENet channel attention mechanism, use BCEloss for confidence loss and category loss, and use CIoUloss for positioning loss;

[0011] S5: Use the training set to train the RStargetNet network, use the validation set to evaluate and optimize the RStargetNet network, and finally use the test set to test and obtain the optimal model;

[0012] S6: Use the optimal model to predict the image in the input network, obtain the category name, confidence score and target box coordinate information of each detection target, and display it in the image to complete the detection.

[0013] Further, step S2 converts the label file in the data set into the YOLOv5 format, specifically including:

[0014] The label file format contains one bounding box per line; the label file format is as follows:

[0015] <class_id><x_center><y_center> <width> <height>

[0016] in<class_id> Indicates the category ID;<x_center> represents the x-coordinate of the center of the bounding box;<y_center> Indicates the y coordinate of the center of the bounding box; <width>Indicates the width of the bounding box; <height>Indicates the height of the bounding box.

[0017] Furthermore, the configuration file of the data set in step S3 is 'hrrsd.yaml', and the specific configuration of the configuration file is as follows:

[0018] train:.. / datasets / HRRSD / images / train

[0019] val:.. / datasets / HRRSD / images / val

[0020] Number of classes: 10

[0021] names:['airplane','ship','storage_tank','baseball_diamond','tennis_court','basketball_court','ground_track_field','harbor','bridge','vehicle'].

[0022] Further, the specific steps of constructing the RStargetNet network in step S4 are as follows:

[0023] S41: Use VanillaNet network to build YOLOv5 Backbone. The Stem layer uses 4×4×3×1024 convolution kernel with a step size of 4 to map the 3-channel remote sensing image to a 1024-channel feature map. A BN layer and ReLU activation function are added after the convolution layer. Stage 1-4 all use MaxPooling, using a 2×2 pooling kernel with a step size of 2, downsampling, and doubling the number of channels. The convolution layer uses a 1×1 convolution kernel, and a BN layer and ReLU activation function are added after the convolution layer.

[0024] S42: Integrate the SENet channel attention mechanism into the PANet part of the YOLOv5 neck network;

[0025] S43: Use DynamicHead dynamic monitoring head to build the YOLOv5head network;

[0026] S44: Both the confidence loss function and the category loss function use binary cross entropy BCE Loss, where the true label is y; for the confidence loss function, if there is an object at a certain position, the label is 1, otherwise it is 0, and the predicted value is y i , as the confidence score predicted by the model, indicates the probability that the target exists at that location; for the category loss function, if the target belongs to a certain category, the label of the corresponding category is 1; otherwise, it is 0, and the predicted value is y i , as the probability of each category predicted by the model;

[0027] The formula for BCELoss is:

[0028]

[0029] S45: The positioning loss is CIoULoss, and the calculation formula of CIoULoss is:

[0030]

[0031] Where IoU is the intersection over union of the predicted box and the true box; ρ 2 (b,b g t) is the Euclidean distance between the center point of the predicted box and the center point of the real box; c is the diagonal distance of the bounding box; v is a metric that reflects the difference in aspect ratio between the predicted box and the real box. The calculation formula of v is:

[0032]

[0033] α is a balance parameter used to adjust the impact on the total loss. The calculation formula of α is:

[0034]

[0035] S46: Combine the backbone network, Neck network, Head network and loss function to obtain the RStargetNet network model.

[0036] Further, step S5 specifically includes:

[0037] S51: Set the input size of the graph to 640×640, use Adam as the optimization function of the network, the initial learning rate is 0.001, use StepLR to adjust the learning rate, the number of epochs for batch training is 300, and the batch size is set to 32; load the pre-trained weights of YOLOv5 yolov5s.pt;

[0038] S52: Set the image input size of the validation set to 640×640. After each epoch, evaluate the validation set and adjust the learning rate to optimize the model performance.

[0039] S53: After each epoch is completed and evaluated on the validation set, the network model parameters are trained and updated, and best.pt is updated, epoch = epoch + 1;

[0040] S54: Repeat steps S51-53 until the RStargetNet network model training is completed, save the optimal model best.pt, and save the model weight last.pt;

[0041] S55: The graphics size of the test set is set to 640×640 and input into the RStargeNet network to normalize the input data.

[0042] The beneficial effects of the present invention are as follows: the RStargetNet network constructed by the present invention improves the detection accuracy and detection speed, and is suitable for detecting small targets such as cars, ships, and airplanes in remote sensing images. In the present invention, the implementation and debugging speed of VanillaNet is relatively fast. Different hyperparameter settings can be quickly iterated and tested, which is suitable for prototype design and preliminary verification. The SENet channel attention mechanism referenced in the model is highly versatile, can improve the generalization ability of the model, and improve the accuracy of small target detection.

[0043] Other advantages, objectives and features of the present invention will be described in the following description to some extent, and to some extent, will be obvious to those skilled in the art based on the following examination and study, or can be taught from the practice of the present invention. The objectives and other advantages of the present invention can be realized and obtained through the following description. BRIEF DESCRIPTION OF THE DRAWINGS

[0044] In order to make the purpose, technical solutions and advantages of the present invention more clear, the present invention will be described in detail below in conjunction with the accompanying drawings, wherein:

[0045] Figure 1 It is a flowchart of the remote sensing small target detection method based on YOLOv5 described in the present invention;

[0046] Figure 2 A flowchart of building the RStargetNet network according to the present invention;

[0047] Figure 3 A flowchart of using the HRRSD data set to train the RStargetNet network according to the present invention;

[0048] Figure 4 This is a flowchart of using RStargetNet to test the HRRSD test set according to the present invention. DETAILED DESCRIPTION

[0049] The following describes the embodiments of the present invention by specific examples, and those skilled in the art can easily understand other advantages and effects of the present invention from the contents disclosed in this specification. The present invention can also be implemented or applied through other different specific embodiments, and the details in this specification can also be modified or changed in various ways based on different viewpoints and applications without departing from the spirit of the present invention. It should be noted that the illustrations provided in the following embodiments only illustrate the basic concept of the present invention in a schematic manner, and the following embodiments and features in the embodiments can be combined with each other without conflict.

[0050] It should be noted that the illustrations provided in the following embodiments are only schematic illustrations of the basic concept of the present invention, and thus the drawings only show components related to the present invention rather than being drawn according to the number, shape and size of components in actual implementation. In actual implementation, the type, quantity and proportion of each component may be changed arbitrarily, and the component layout may also be more complicated.

[0051] In the following description, numerous details are discussed to provide a more thorough explanation of the embodiments of the present invention. However, it is obvious to those skilled in the art that the embodiments of the present invention can be implemented without these specific details. In other embodiments, well-known structures and devices are shown in the form of block diagrams rather than in detail to avoid making the embodiments of the present invention difficult to understand.

[0052] like Figure 1 As shown, the present invention provides a remote sensing small target detection method based on YOLOv5, the steps are as follows:

[0053] Step 1), collect data sets. The HRRSD data set used in this embodiment contains 21761 images. There are 55740 target instances in HRRSD, and each category is about 4k. HRRSD contains 13 types of targets. The 13 categories are: airplanes, baseball fields, basketball courts, bridges, intersections, athletic fields, ports, parking lots, ships, storage tanks, T-junctions, tennis courts, and cars. The sample size between each category is relatively balanced, and each category has about 4000 samples. The data set is divided into several subsets, and the number of images in each subset is 5401 for 'train', 5417 for 'val', and 10943 for 'test'. The 'train-val' subset is the merger of 'train' and 'val'.

[0054] Step 2) Preprocess the text data and divide the data into train.txt and test.txt for separate processing;

[0055] Step 3) Convert the label file in the HRRSD dataset into a ".txt" file in YOLOv5 format

[0056] The label file format contains one bounding box per line.

[0057] The format is as follows:<class_id><x_center><y_center> <width> <height>

[0058] The information represented by each formula is

[0059] <class_id> : Category ID (starting from 0)

[0060] <x_center> : The x-coordinate of the center of the bounding box (normalized to 0-1)

[0061] <y_center> : The y coordinate of the center of the bounding box (normalized to 0-1)

[0062] <width>: The width of the bounding box (normalized to 0-1)

[0063] <height>: The height of the bounding box (normalized to 0-1)

[0064] Step 4) Create the dataset configuration file 'hrrsd.yaml'

[0065] The specific configuration of the .yaml file is as follows:

[0066] train:.. / datasets / HRRSD / images / train

[0067] val:.. / datasets / HRRSD / images / val

[0068] Number of classes: 10

[0069] names:['airplane','ship','storage_tank','baseball_diamond','tennis_court','basketball_court','ground_track_field','harbor','bridge','vehicle']

[0070] Step 5) Build RStargetNet network: Figure 2 As shown in the figure, based on the YOLOv5 network, VanillaNet is used to replace the CSPDarknet53 backbone network, and the SENet channel attention mechanism is added. The confidence loss, category loss and positioning loss use BCEloss, BCEloss and CIoUloss respectively.

[0071] Step 51), when using the VanillaNet network to build the Backbone of YOLOv5, the Stem layer uses a 4×4×3×1024 convolution kernel with a step size of 4 to map the 3-channel remote sensing image to a 1024-channel feature map. A BN layer and a ReLU activation function are added after the convolution layer. Stage 1-4 all use MaxPooling, using a 2×2 pooling kernel with a step size of 2, downsampling, and doubling the number of channels. The convolution layer uses a 1×1 convolution kernel, and a BN layer and a ReLU activation function are added after the convolution layer.

[0072] Step 52) Integrate the SENet channel attention mechanism into the PANet part of the YOLOv5 neck network to enhance the overall network's ability to express features.

[0073] Step 53), use DynamicHead dynamic monitoring head to build the YOLOv5head network;

[0074] Step 54), the confidence loss function uses binary cross entropy BCE Loss, where the true label is y: if there is an object at a certain position, the label is 1; otherwise it is 0. The predicted value is y i , as the confidence score predicted by the model, represents the probability that a target exists at that location. The formula for BCELoss is:

[0075]

[0076] Step 55), the category loss function uses binary cross entropy BCE Loss, where the true label is y: if the target belongs to a certain category, the label of the corresponding category is 1; otherwise it is 0. The predicted value is y i , as the probability of each category predicted by the model. The formula of BCELoss is:

[0077]

[0078] Step 56), the positioning loss is CIoULoss, CIoU (Complete Intersection over Union) Loss is an improved IoU (Intersection over Union) loss function, which is specifically used for bounding box regression in target detection tasks. It introduces considerations of bounding box shape and center point distance on the basis of standard IoU, thereby providing a more accurate loss metric. IoU is the intersection over union ratio of the predicted box and the true box, ρ 2 (b,b g t) is the Euclidean distance between the center point of the predicted box and the center point of the true box, and c is the diagonal distance of the bounding box. v is a metric that reflects the difference in aspect ratio between the predicted box and the true box. α is a balancing parameter used to adjust the impact on the total loss.

[0079] The calculation formula for v is:

[0080]

[0081] The calculation formula of α is:

[0082]

[0083] The calculation formula of CIoULoss is:

[0084]

[0085] The backbone network, Neck network, Head network and loss function are combined to obtain the RStargetNet network model.

[0086] Step 6), such as Figure 3 As shown, the RStargetNet network is trained using the training set in the HRRSD dataset;

[0087] Set the input size of the image to 640×640;

[0088] Adam is used as the optimization function of the network, the initial learning rate is 0.001, and StepLR is used to adjust the learning rate. The number of epochs for batch training is 300, and the batch size is set to 32;

[0089] The model loads the pre-trained weights of YOLOv5 yolov5s.pt;

[0090] The image input size of the validation set is set to 640×640. After each epoch, the model is evaluated on the validation set and the learning rate is adjusted to optimize the model performance.

[0091] After each epoch is completed and evaluated on the validation set, the network model parameters are updated and best.pt is updated, epoch = epoch + 1;

[0092] Repeat the operation until the RStargetNet network model training is completed, save the optimal model best.pt, and save the model weight last.pt;

[0093] like Figure 4 As shown, the graphics size of the test set is set to 640×640 and input into the RStargeNet network;

[0094] Normalize the input data;

[0095] Step 7) Load the optimal model to predict the image in the input network;

[0096] Obtain the category name, confidence score, and corresponding target box coordinate information of each detected target and display them in the picture to complete the detection.

[0097] In the above embodiments, the description's reference to "this embodiment" indicates that a particular feature, structure, or characteristic described in conjunction with the embodiment is included in at least some embodiments, but not necessarily all embodiments. Multiple occurrences of "this embodiment" do not necessarily all refer to the same embodiment.

[0098] In the above-described embodiments, although the invention has been described in conjunction with specific embodiments of the invention, many substitutions, modifications, and variations of these embodiments will be apparent to those of ordinary skill in the art based on the foregoing description. For example, other storage structures (e.g., dynamic RAM (DRAM)) may use the embodiments discussed. Embodiments of the invention are intended to encompass all such substitutions, modifications, and variations that fall within the broad scope of the appended claims.

[0099] This embodiment further provides a computer-readable storage medium on which a computer program is stored. When the program is executed by a processor, any one of the methods in this embodiment is implemented.

[0100] This embodiment also provides an electronic terminal, including: a processor and a memory;

[0101] The memory is used to store a computer program, and the processor is used to execute the computer program stored in the memory, so that the terminal executes any one of the methods in this embodiment.

[0102] The computer-readable storage medium in this embodiment can be understood by ordinary technicians in this field: all or part of the steps of implementing the above-mentioned method embodiments can be completed by hardware related to the computer program. The aforementioned computer program can be stored in a computer-readable storage medium. When the program is executed, the steps of the above-mentioned method embodiments are executed; and the aforementioned storage medium includes: ROM, RAM, magnetic disk or optical disk and other media that can store program codes.

[0103] The electronic terminal provided in this embodiment includes a processor, a memory, a transceiver and a communication interface. The memory and the communication interface are connected to the processor and the transceiver and complete communication with each other. The memory is used to store computer programs, the communication interface is used to communicate, and the processor and the transceiver are used to run computer programs so that the electronic terminal executes each step of the above method.

[0104] In this embodiment, the memory may include a random access memory (RAM), and may also include a non-volatile memory (non-volatile memory), such as at least one disk memory.

[0105] The above-mentioned processor can be a general-purpose processor, including a central processing unit (CPU), a network processor (NP), etc.; it can also be a digital signal processor (DSP), an application specific integrated circuit (ASIC), a field programmable gate array (FPGA) or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components.

[0106] The present invention can be used in many general or special computing system environments or configurations, such as personal computers, server computers, handheld or portable devices, tablet devices, multiprocessor systems, microprocessor-based systems, set-top boxes, programmable consumer electronic devices, network PCs, minicomputers, mainframe computers, distributed computing environments including any of the above systems or devices, and the like.

[0107] The present invention may be described in the general context of computer-executable instructions executed by a computer, such as program modules. Generally, program modules include routines, programs, objects, components, data structures, etc. that perform specific tasks or implement specific abstract data types. The present invention may also be practiced in distributed computing environments where tasks are performed by remote processing devices connected through a communication network. In a distributed computing environment, program modules may be located in local and remote computer storage media, including storage devices.

[0108] Finally, it should be noted that the above embodiments are only used to illustrate the technical solution of the present invention rather than to limit it. Although the present invention has been described in detail with reference to the preferred embodiments, those skilled in the art should understand that the technical solution of the present invention can be modified or replaced by equivalents without departing from the purpose and scope of the technical solution, which should be included in the scope of the claims of the present invention.< / height> < / width> < / height> < / width> < / height> < / width> < / height> < / width>

Claims

1. A remote sensing small target detection method based on YOLOv5, characterized by: The following steps are involved: S1: Collect remote sensing image data sets, preprocess the data sets, and divide them into training sets, validation sets, and test sets; S2: Convert the label file in the dataset to YOLOv5 format; S3: Create a configuration file for the dataset; S4: Build RStargetNet network: Based on the YOLOv5 network, use VanillaNet to replace the CSPDarknet53 backbone network, add SENet channel attention mechanism, use BCEloss for confidence loss and category loss, and use CIoUloss for positioning loss; S5: Use the training set to train the RStargetNet network, use the validation set to evaluate and optimize the RStargetNet network, and finally use the test set to test and obtain the optimal model; S6: Use the optimal model to predict the image in the input network, obtain the category name, confidence score and target box coordinate information of each detection target, and display it in the image to complete the detection.

2. The remote sensing small target detection method based on YOLOv5 according to claim 1, characterized in that: Step S2 converts the label file in the data set into the YOLOv5 format, specifically including: The label file format contains one bounding box per line; the label file format is as follows: <class_id><x_center><y_center> <width> <height>< / height> < / width> in<class_id> Indicates the category ID;<x_center> represents the x-coordinate of the center of the bounding box;<y_center> Indicates the y coordinate of the center of the bounding box; <width>Indicates the width of the bounding box; <height> Indicates the height of the bounding box.< / height> < / width> 3. The remote sensing small target detection method based on YOLOv5 according to claim 1, characterized in that: The configuration file of the data set in step S3 is 'hrrsd.yaml', and the specific configuration of the configuration file is as follows: train:.. / datasets / HRRSD / images / train val:.. / datasets / HRRSD / images / val Number of classes: 10 names:['airplane','ship','storage_tank','baseball_diamond','tennis_court','basketball_court','ground_track_field','harbor','bridge','vehicle'].

4. The remote sensing small target detection method based on YOLOv5 according to claim 1, characterized in that: The specific steps of constructing the RStargetNet network in step S4 are as follows: S41: Use VanillaNet network to build YOLOv5 Backbone. The Stem layer uses 4×4×3×1024 convolution kernel with a step size of 4 to map the 3-channel remote sensing image to a 1024-channel feature map. A BN layer and ReLU activation function are added after the convolution layer. Stage 1-4 all use MaxPooling, using a 2×2 pooling kernel with a step size of 2, downsampling, and doubling the number of channels. The convolution layer uses a 1×1 convolution kernel, and a BN layer and ReLU activation function are added after the convolution layer. S42: Integrate the SENet channel attention mechanism into the PANet part of the YOLOv5 neck network; S43: Use DynamicHead dynamic monitoring head to build the YOLOv5head network; S44: Both the confidence loss function and the category loss function use binary cross entropy BCE Loss, where the true label is y; for the confidence loss function, if there is an object at a certain position, the label is 1, otherwise it is 0, and the predicted value is y i , as the confidence score predicted by the model, indicates the probability that the target exists at that location; for the category loss function, if the target belongs to a certain category, the label of the corresponding category is 1; otherwise, it is 0, and the predicted value is y i , as the probability of each category predicted by the model; The formula for BCELoss is: S45: The positioning loss is CIoULoss, and the calculation formula of CIoULoss is: Where IoU is the intersection over union of the predicted box and the true box; ρ 2 (b,b g t) is the Euclidean distance between the center point of the predicted box and the center point of the real box; c is the diagonal distance of the bounding box; v is a metric that reflects the difference in aspect ratio between the predicted box and the real box. The calculation formula of v is: α is a balance parameter used to adjust the impact on the total loss. The calculation formula of α is: S46: Combine the backbone network, Neck network, Head network and loss function to obtain the RStargetNet network model.

5. The remote sensing small target detection method based on YOLOv5 according to claim 1, characterized in that: Step S5 specifically includes: S51: Set the input size of the graph to 640×640, use Adam as the optimization function of the network, the initial learning rate is 0.001, use StepLR to adjust the learning rate, the number of epochs for batch training is 300, and the batch size is set to 32; load the pre-trained weights of YOLOv5 yolov5s.pt; S52: Set the image input size of the validation set to 640×640. After each epoch, evaluate the validation set and adjust the learning rate to optimize the model performance. S53: After each epoch is completed and evaluated on the validation set, the network model parameters are trained and updated, and best.pt is updated, epoch = epoch + 1; S54: Repeat steps S51-53 until the RStargetNet network model training is completed, save the optimal model best.pt, and save the model weight last.pt; S55: The graphics size of the test set is set to 640×640 and input into the RStargeNet network to normalize the input data.

Images