Small target detection method and system based on dual-network cascade model

Abstract

The invention relates to a small target detection method based on a dual-network cascade model. The method comprises the following steps: obtaining an image to be detected; constructing a target detection model to detect a target object in the to-be-detected image to obtain a target detection result; the target detection model is constructed based on a dual-network cascade model, and comprises a first target detection network used for detecting a target associated object in the to-be-detected image to obtain a target bounding box; and the second target detection network is constructed based on an improved YOLO-NAS network and is used for determining a target associated image according to the target bounding box and further detecting a target object in the target associated image to obtain a target detection result. According to the invention, the detection accuracy and precision of the small target can be improved, and the small target in a picture, especially a large picture, can be effectively detected.

Description

Technical Field

[0001] The present invention relates to the field of computer image vision technology, and in particular to a small target detection method and system based on a dual-network cascade model. Background Art

[0002] At present, in terms of detection steps, target detection algorithms can be divided into two categories, one is the two-stage algorithm and the other is the one-stage algorithm. The advantage of the one-stage algorithm is its fast speed, but its detection accuracy is relatively low. The two-stage network mainly completes the multi-level classification problem and does not involve secondary refinement of target recognition. At the same time, small object detection (SOD) has been a hot and difficult problem in recent years and is one of the key technologies for target detection. Although target detection has achieved great success, there are still difficulties in detecting small targets. Since small targets occupy a very small area of ​​the image and are usually long-distance targets, there are few target pixels that can be detected, it is difficult to capture effective features, and they are easily disturbed by the background. The sensitivity of the neural network to target features makes its detection effect and large targets have a significant performance gap. However, early target detection was mainly aimed at conventional size targets, and small targets have only received attention in recent years. Therefore, small target detection is mostly based on models of conventional size detection. As a subfield of target detection, it mainly studies the detection of targets with smaller sizes. It has important theoretical and practical significance in scenes such as ground detection taken by drones, cell detection in medical images, monitoring, and pedestrian detection. When extracting features with the help of convolution and pooling operations, the feature information of small targets will gradually decrease or even be lost as the number of network layers increases, making it difficult for the network to detect them. In particular, when detecting small targets in large images containing many objects, there are problems such as insufficient model detection accuracy and low target detection accuracy due to excessive image information. Therefore, a target detection method is needed that can solve the problem of insufficient effective features for small targets, thereby improving the detection accuracy of small targets, especially the detection accuracy of small targets in large images. Summary of the invention

[0003] The technical problem to be solved by the present invention is to provide a small target detection method and system based on a dual network cascade model, which can improve the small target detection accuracy and effectively detect small targets in pictures, especially large pictures.

[0004] The technical solution adopted by the present invention to solve the technical problem is: to provide a small target detection method based on a dual network cascade model, comprising the following steps:

[0005] Acquire the image to be tested;

[0006] Constructing a target detection model to detect the target object in the image to be tested to obtain a target detection result; the target detection model is constructed based on a dual network cascade model, including:

[0007] A first target detection network is used to detect target-related objects in the image to be detected to obtain a target bounding box;

[0008] The second target detection network is constructed based on the improved YOLO-NAS network, and is used to determine the target associated image according to the target bounding box, and then detect the target object in the target associated image to obtain the target detection result.

[0009] Furthermore, the improved YOLO-NAS network includes:

[0010] Introduce the channel attention mechanism with embedded position information in the feature extraction network to extract image features;

[0011] The multi-scale feature fusion module is used to perform feature fusion and provide input feature map representation for the two detection heads with different receptive fields.

[0012] Optimize the regression loss function.

[0013] Furthermore, the regression loss function L of the improved YOLO-NAS network is:

[0014]

[0015] Among them, b, w, h represent the prediction box center, prediction box width and prediction box height respectively, b gt , w gt ,h gt They represent the center of the real box, the width of the real box and the height of the real box respectively, c w and c h Represents the width and height of the minimum bounding box of the real box and the predicted box respectively.

[0016] Furthermore, the detecting the target-related object in the image to be detected to obtain a target bounding box includes:

[0017] Analyzing the training image samples to obtain the target-associated object and a plurality of first position prediction results corresponding to the target-associated object and their confidences;

[0018] The most relevant first position prediction result is extracted as a target bounding box according to the confidence of the first position prediction result.

[0019] Further, extracting the most relevant first position prediction result according to the confidence of the first position prediction result includes:

[0020] Extracting the first position prediction result having a confidence level greater than a confidence level threshold;

[0021] The extracted first position prediction result is subjected to non-maximum suppression processing to obtain the most relevant first position prediction result.

[0022] Further, determining the target associated image according to the target bounding box includes:

[0023] Setting the coordinates of the upper left corner of the target associated image to a first preset multiple of the coordinates of the upper left corner of the target bounding box;

[0024] Setting the coordinates of the lower right corner of the target associated image to a second preset multiple of the coordinates of the lower right corner of the target bounding box;

[0025] The first preset multiple is smaller than the second preset multiple.

[0026] Furthermore, the target detection model is constructed by the following method:

[0027] Obtain several training image samples;

[0028] Build an object detection model based on the number of preset detection categories;

[0029] Using a rectangular frame to mark the category label and position label of the preset detection category in each of the training image samples;

[0030] Analyzing the training image samples using the target detection model to obtain prediction results;

[0031] The prediction result is compared with the category label and the position label, and network parameters of the first object detection network and the second object detection network are optimized according to the comparison result.

[0032] Furthermore, the first target detection network and the second target detection network are both built based on the YOLO-NAS network.

[0033] The present invention also provides a small target detection system based on a dual network cascade model, comprising:

[0034] Input module, used to obtain the image to be tested;

[0035] The analysis module is used to construct a target detection model to detect the target object in the image to be tested and obtain a target detection result; the target detection model is constructed based on a dual network cascade model, including:

[0036] A first target detection network is used to detect target-related objects in the image to be detected to obtain a target bounding box;

[0037] The second target detection network is constructed based on the improved YOLO-NAS network, and is used to determine the target associated image according to the target bounding box, and then detect the target object in the target associated image to obtain the target detection result.

[0038] Furthermore, the improved YOLO-NAS network includes:

[0039] A feature extraction part, used to extract image features from the target associated image, and introduce a channel attention mechanism that embeds position information in the process of image feature extraction;

[0040] It is used to perform feature fusion through a multi-scale feature fusion module to provide input feature map representation for two detection heads with different receptive fields.

[0041] Furthermore, the regression loss function L of the improved YOLO-NAS network is:

[0042]

[0043] Among them, b, w, h represent the prediction box center, prediction box width and prediction box height respectively, b gt , w gt ,h gt They represent the center of the real box, the width of the real box and the height of the real box respectively, c w and c h Represents the width and height of the minimum bounding box of the real box and the predicted box respectively.

[0044] Beneficial Effects

[0045] Due to the adoption of the above technical solution, the present invention has the following advantages and positive effects compared with the prior art: the present invention constructs a target detection model by cascading two YOLO-NAS networks, uses the target bounding box image obtained by the first target detection network as the input image of the second target detection network, and outputs the final detection result after analysis by the second target detection network, and improves the network structure of the second YOLO-NAS network for the dual network cascade framework, including embedding a CA attention mechanism, deleting a large-scale detection head, and optimizing a regression loss function, so that it is more focused on small target detection, improving the accuracy and precision of small target detection, and the structure of the two network cascades allows more flexible network tuning, further improving the overall performance; wherein, embedding a CA attention mechanism can focus on a larger area without introducing significant overhead, thereby improving the detection capability of small targets; deleting a large-scale detection head can save computing power waste, and can also improve detection speed and detection performance; the optimized regression loss function takes into account the overlapping area, the distance of the center point, and the true difference of height and width, which can improve its convergence speed and better positioning. BRIEF DESCRIPTION OF THE DRAWINGS

[0046] Figure 1 is a flow chart of a first embodiment of the present invention;

[0047] Figure 2 is a schematic diagram of a first embodiment of the present invention;

[0048] Figure 3 It is a schematic diagram of the second-level network model architecture of the first embodiment of the present invention. DETAILED DESCRIPTION

[0049] The present invention will be further described below in conjunction with specific embodiments. It should be understood that these embodiments are only used to illustrate the present invention and are not intended to limit the scope of the present invention. In addition, it should be understood that after reading the content taught by the present invention, those skilled in the art can make various changes or modifications to the present invention, and these equivalent forms fall within the scope limited by the appended claims of the application equally.

[0050] The first embodiment of the present invention relates to a small target dual-network cascade detection method based on a YOLO network, comprising the following steps:

[0051] Step S1: Model training: mainly completes two annotations of the same data set, including the annotation of the target object and the target-related objects, and puts the data set and the annotation results into the model for training; the specific steps are as follows:

[0052] Step S1.1: After obtaining the image set, use a rectangular box to mark the category and location coordinates of each image;

[0053] Step S1.2: Use the YOLO_NAS model to train the image to obtain the training results of the first-level network target;

[0054] Step S1.3: Use the improved YOLO_NAS model training image to obtain the training results of the second-level network target; specifically including:

[0055] Step S1.3.1: Optimization of YOLO_NAS model architecture, such as Figure 3As shown in the figure. Since small targets are easily lost or easily disturbed by other backgrounds during feature extraction, the network is considered to pay more attention to the channels that have the ability to recognize small targets by learning the correlation between channels. The CA attention mechanism is introduced to help the model more accurately locate and identify targets of interest, realize feature selection, and improve generalization ability. Since the second target detection network needs to pay more attention to small targets, and the network's detection heads are respectively for large-scale, medium-scale, and small-scale targets, for the second target detection network, the practical application significance of the detection head for detecting large-scale targets will be greatly reduced, while increasing the amount of calculation and reasoning speed. Therefore, the second target detection network model will delete the detection head for large-scale detection on the basis of the original model to improve the detection of small targets.

[0056] Step S1.3.2: In the original loss function, when the predicted box and the true box overlap, that is, when the predicted box is inside the true box, the original loss function will not be able to judge the positional relationship between the predicted box and the true box, resulting in slower regression convergence. Therefore, the center point prediction method, that is, the distance loss, is considered. By judging the relationship between the square of the center distance between the predicted box and the true box and the two boxes, the target box regression becomes more stable. In addition, since the original loss function does not take into account the length and width issues, the shape of the predicted bounding box will be inconsistent with the true bounding box. Therefore, the proportional relationship between width and height is considered, and the aspect ratio of the predicted box and the true box is introduced, so that the loss function pays more attention to the shape of the bounding box. The optimized loss function comprehensively considers the overlapping area loss, width and height loss, and center distance loss. Among them, b, w, and h represent the center of the prediction box, the width of the prediction box, and the height of the prediction box, respectively, and b gt , w gt ,h gt They are the center of the real box, the width of the real box and the height of the real box, c w and c h Represents the width and height of the minimum bounding box of the real box and the predicted box respectively.

[0057]

[0058] Step S1.4: Generate the model weights of the first-level network and the second-level network of the dual-network cascade network through steps S1.2 and S1.3, and store them in a file in pth format.

[0059] Step S2: Figure 1 As shown, the first target detection network detects the range coordinates of the target-related objects, and the second target detection network further detects the relevant pictures cropped from the original picture according to the coordinates of the first-level network. The specific steps are as follows:

[0060] Step S2.1: Define the coordinate array detections_xyxy[], the label array detections_labels[], and the confidence array detections_confidence[];

[0061] Step S2.2: Prepare test data: traverse the test data, read the i-th data, that is, the i-th picture, use the first object detection network to predict the picture by loading the trained model weights, and obtain data including predicted coordinates, labels, and confidence levels; specifically, include:

[0062] Step S2.2.1: Obtain the number of categories through the defined data parameters and build the YOLO-NAS model;

[0063] Step S2.2.2: The image input to the model is scaled to 640*640, and the prediction result is obtained through the forward propagation of the model;

[0064] Step S2.2.3: Post-process the output of the model to obtain the prediction results of bounding boxes and category labels. For each category, perform threshold filtering based on the confidence level to remove the prediction results with low probability. For the retained prediction results, perform non-maximum suppression processing to eliminate overlapping bounding boxes and retain only the most relevant bounding boxes;

[0065] Step S2.2.4: Target box drawing: Draw the target on the original image to show the target object and location category detected by the model.

[0066] Step S2.3: Traverse each predicted result, select the predicted coordinate information, and intercept the input image of the second target detection network through the target coordinates. Since it is impossible to ensure that the detection frame obtained by the first target detection network is completely consistent with expectations, in order to reduce the missed detection rate, it is necessary to expand the range of the detection frame. In this embodiment, the upper left corner coordinate (x_min, y_min) is set to 0.97 times the predicted result, and the lower right corner coordinate (x_max, y_max) is set to 1.02 times the predicted result;

[0067] Step S2.4: The input prediction model is the image captured in S2.3. Use the second object detection network to predict the image by loading the trained model weights to obtain data including prediction coordinates, labels, and confidence levels;

[0068] Step S2.5: Add the predicted coordinate information of an image to the detections_xyxy array, the label information to the detections_labels array, and the confidence array to the detections_confidence array. Calculate the real label information and the predicted information to obtain the evaluation index of the model.

[0069] like Figure 2 As shown in the figure, taking the detection of whether a helmet is worn on an electric bike as an example, in real scenarios, the premise of detecting whether a helmet is worn (target) is riding an electric bike. Therefore, the first target detection network detects the riding electric bike, which is an object closely related to the final target. When the first target detection network detects the target, the detection result is transmitted to the second target detection network, which detects the final target and obtains the final position.

[0070] The second embodiment of the present invention relates to a small target dual-network cascade detection system based on a YOLO network, comprising:

[0071] Input module, used to obtain the image to be tested;

[0072] An analysis module is used to construct a target detection model to detect the target object in the image to be tested and obtain the target detection result;

[0073] The target detection model is built based on a dual network cascade model, including:

[0074] The first target detection network is built based on the YOLO-NAS network and is used to detect target-related objects in the image to be tested and obtain the target bounding box;

[0075] The second target detection network is built based on the improved YOLO-NAS network, which is used to determine the target-associated image according to the target bounding box, and then detect the target object in the target-associated image to obtain the target detection result.

[0076] When training a model, the same data set is annotated twice, and then the data set and the annotated results are put into the model for training. The specific steps are as follows:

[0077] Step S1.1: After obtaining the image set, use a rectangular box to mark the category and location coordinates of each image;

[0078] Step S1.2: Use the YOLO_NAS model to train the image to obtain the training results of the first-level network target;

[0079] Step S1.3: Design an optimized model training picture based on the YOLO_NAS model to obtain the training results of the second-level network target; specifically including:

[0080] Step S1.3.1: Optimization of YOLO_NAS model architecture, such as Figure 3As shown in the figure. Since small targets are easily lost or easily disturbed by other backgrounds during feature extraction, the network is considered to pay more attention to the channels that have the ability to recognize small targets by learning the correlation between channels. The CA attention mechanism is introduced to help the model more accurately locate and identify targets of interest, realize feature selection, and improve generalization ability. Since the second target detection network needs to pay more attention to small targets, and the network's detection heads are respectively for large-scale, medium-scale, and small-scale targets, for the second target detection network, the practical application significance of the detection head for detecting large-scale targets will be greatly reduced, while increasing the amount of calculation and reasoning speed. Therefore, the second target detection network model will delete the detection head for large-scale detection on the basis of the original model to improve the detection of small targets.

[0081] Step S1.3.2: In the original loss function, when the predicted box and the true box overlap, that is, when the predicted box is inside the true box, the original loss function will not be able to judge the positional relationship between the predicted box and the true box, resulting in slower regression convergence. Therefore, the center point prediction method, that is, the distance loss, is considered. By judging the relationship between the square of the center distance between the predicted box and the true box and the two boxes, the target box regression becomes more stable. In addition, since the original loss function does not take into account the length and width issues, the shape of the predicted bounding box will be inconsistent with the true bounding box. Therefore, the proportional relationship between width and height is considered, and the aspect ratio of the predicted box and the true box is introduced, so that the loss function pays more attention to the shape of the bounding box. The optimized loss function comprehensively considers the overlapping area loss, width and height loss, and center distance loss. Among them, b, w, and h represent the center of the prediction box, the width of the prediction box, and the height of the prediction box, respectively, and b gt , w gt ,h gt They are the center of the real box, the width of the real box and the height of the real box, c w and c h Represents the width and height of the minimum bounding box of the real box and the predicted box respectively.

[0082]

[0083] Step S1.4: Generate the model weights of the first-level network and the second-level network of the dual-network cascade network through steps S1.2 and S1.3, and store them in a file in pth format.

[0084] During detection, the first target detection network detects the range coordinates of the target-related objects, and the second target detection network further detects the relevant pictures cropped from the original picture according to the coordinates of the first-level network. The specific steps are as follows:

[0085] Step S2.1: Define the coordinate array detections_xyxy[], the label array detections_labels[], and the confidence array detections_confidence[];

[0086] Step S2.2: Prepare test data: traverse the test data, read the i-th data, that is, the i-th picture, use the first object detection network to predict the picture by loading the trained model weights, and obtain data including predicted coordinates, labels, and confidence levels; specifically, include:

[0087] Step S2.2.1: Obtain the number of categories through the defined data parameters and build the YOLO-NAS model;

[0088] Step S2.2.2: The image input to the model is scaled to 640*640, and the prediction result is obtained through the forward propagation of the model;

[0089] Step S2.2.3: Post-process the output of the model to obtain the prediction results of bounding boxes and category labels. For each category, perform threshold filtering based on the confidence level to remove the prediction results with low probability. For the retained prediction results, perform non-maximum suppression processing to eliminate overlapping bounding boxes and retain only the most relevant bounding boxes;

[0090] Step S2.2.4: Target box drawing: Draw the target on the original image to show the target object and location category detected by the model.

[0091] Step S2.3: Traverse each predicted result, select the predicted coordinate information, and intercept the input image of the second target detection network through the target coordinates. Since it is impossible to ensure that the detection frame obtained by the first target detection network is completely consistent with expectations, in order to reduce the missed detection rate, it is necessary to expand the range of the detection frame. In this embodiment, the upper left corner coordinate (x_min, y_min) is set to 0.97 times the predicted result, and the lower right corner coordinate (x_max, y_max) is set to 1.02 times the predicted result;

[0092] Step S2.4: The input prediction model is the image captured in S2.3. Use the second object detection network to predict the image by loading the trained model weights to obtain data including prediction coordinates, labels, and confidence levels;

[0093] Step S2.5: Add the predicted coordinate information of an image to the detections_xyxy array, the label information to the detections_labels array, and the confidence array to the detections_confidence array. Calculate the real label information and the predicted information to obtain the evaluation index of the model.

Claims

1. A small target detection method based on a dual network cascade model, characterized in that: The following steps are involved: Acquire the image to be tested; Constructing a target detection model to detect the target object in the image to be tested to obtain a target detection result; The target detection model is built based on a dual network cascade model, including: A first target detection network is used to detect target-related objects in the image to be detected to obtain a target bounding box; The second target detection network is constructed based on the improved YOLO-NAS network, and is used to determine the target associated image according to the target bounding box, and then detect the target object in the target associated image to obtain the target detection result.

2. The method according to claim 1, characterized in that The improved YOLO-NAS network includes: Introduce a channel attention mechanism with embedded position information in the feature extraction network to enhance image feature extraction; The multi-scale feature fusion module is used to perform feature fusion and provide input feature map representation for the two detection heads with different receptive fields. Optimize the regression loss function.

3. The system according to claim 2, characterized in that The regression loss function L of the improved YOLO-NAS network is: Among them, b, w, h represent the prediction box center, prediction box width and prediction box height respectively, b gt , w gt ,h gt They represent the center of the real box, the width of the real box and the height of the real box respectively, c w and c h Represents the width and height of the minimum bounding box of the real box and the predicted box respectively.

4. The method according to claim 1, characterized in that The detecting the target-related object in the image to be detected to obtain a target bounding box includes: Analyzing the training image samples to obtain the target-associated object and a plurality of first position prediction results corresponding to the target-associated object and their confidences; The most relevant first position prediction result is extracted as a target bounding box according to the confidence of the first position prediction result.

5. The method according to claim 4, characterized in that The extracting the most relevant first position prediction result according to the confidence of the first position prediction result includes: Extracting the first position prediction result having a confidence level greater than a confidence level threshold; The extracted first position prediction result is subjected to non-maximum suppression processing to obtain the most relevant first position prediction result.

6. The method according to claim 1, characterized in that Determining the target associated image according to the target bounding box comprises: Setting the coordinates of the upper left corner of the target associated image to a first preset multiple of the coordinates of the upper left corner of the target bounding box; Setting the coordinates of the lower right corner of the target associated image to a second preset multiple of the coordinates of the lower right corner of the target bounding box; The first preset multiple is smaller than the second preset multiple.

7. The method according to claim 1, characterized in that The target detection model is constructed by the following method: Obtain several training image samples; Build an object detection model based on the number of preset detection categories; Using a rectangular frame to mark the category label and position label of the preset detection category in each of the training image samples; Analyzing the training image samples using the target detection model to obtain prediction results; The prediction result is compared with the category label and the position label, and network parameters of the first object detection network and the second object detection network are optimized according to the comparison result.

8. A small target detection system based on a dual network cascade model, characterized in that: include: Input module, used to obtain the image to be tested; An analysis module is used to construct a target detection model to detect the target object in the image to be tested and obtain a target detection result; The target detection model is built based on a dual network cascade model, including: A first target detection network is used to detect target-related objects in the image to be detected to obtain a target bounding box; The second target detection network is constructed based on the improved YOLO-NAS network, and is used to determine the target associated image according to the target bounding box, and then detect the target object in the target associated image to obtain the target detection result.

9. The system according to claim 8, characterized in that The improved YOLO-NAS network includes: A feature extraction network is used to enhance image feature extraction of the target associated image and introduce a channel attention mechanism that embeds position information in the process of image feature extraction; It is used to perform feature fusion through a multi-scale feature fusion module to provide input feature map representation for two detection heads with different receptive fields. Optimize the regression loss function.

10. The system according to claim 9, characterized in that The regression loss function L of the improved YOLO-NAS network is: Among them, b, w, h represent the prediction box center, prediction box width and prediction box height respectively, b gt , w gt ,h gt They represent the center of the real box, the width of the real box and the height of the real box respectively, c w and c h Represents the width and height of the minimum bounding box of the real box and the predicted box respectively.

Images