Small target detection network based on wavelet convolution enhanced YOLOv8

Abstract

The invention relates to the field of computer vision and target detection, in particular to the field of rotating small target detection based on a convolutional neural network, and particularly relates to a method for improving a YOLOv8 model by introducing a Haar wavelet transform down-sampling module HWD, a wavelet transform feature enhancement module WTFEM and a bounding box regression loss function MPDIOU based on the minimum point distance. Therefore, the detection precision and robustness of the small target are improved. The method comprises the following steps: firstly, replacing a down-sampling module in YOLOv8 with a Haar wavelet transform down-sampling module HWD; the HWD uses Haar wavelet transform to reduce the spatial resolution of the feature map, and at the same time, more information is reserved as much as possible. And secondly, a Wavelet Transform Feature Enhancement Module (WTFEM) is innovatively introduced into a check part of the network, so that the limitation of a traditional feature fusion mode is broken through, and the semantic understanding and detail retention capability of the model on a multi-scale target is remarkably improved. And finally, replacing the original loss function with a bounding box regression loss function MPDIOU based on the minimum point distance. According to the loss function, the Euclidean distance of the nearest vertex between a prediction frame and a real frame is calculated, and an area overlapping rate and a central point distance optimization target are combined, so that the problem of gradient disappearance of a traditional IoU in a boundary frame non-overlapping scene is solved.

Description

Technical Field

[0001] The present invention relates to the fields of computer vision and target detection, in particular to the field of rotating small target detection based on convolutional neural networks. Specifically, the YOLOv8 model is improved by introducing a Haar wavelet downsampling module (HWD), a wavelet transform feature enhancement module (WTFEM), and a bounding box regression loss function based on minimum point distance (MPDIoU) to achieve high-precision detection of small targets in remote sensing images. Background Art

[0002] YOLOv8 is a next-generation target detection framework launched by Ultralytics. It inherits the core advantages of the YOLO series of single-stage detectors and achieves significant breakthroughs in speed and accuracy. Its architecture adopts a modular design. The backbone network (Backbone) is composed of an improved CSPDarknet53. The C2F module replaces the traditional C3 structure to enhance gradient flow and reduce parameter redundancy. At the same time, the SPPF layer is introduced to enhance multi-scale feature extraction capabilities. The detection head adopts a decoupled-head structure to separate the classification and regression tasks. For the first time, the Anchor-Free mechanism is introduced to directly predict the target center point, avoiding the complexity of anchor box size adaptation in traditional anchor-based methods.

[0003] In terms of training strategy, YOLOv8 innovatively adopts the Task-Aligned Assigner dynamic label assignment algorithm, combined with Distribution Focal Loss to optimize small object detection performance, and significantly improves model generalization capabilities through mixed-precision training and distributed data augmentation (such as Mosaic enhancement and Copy-Paste technology). Experimental data shows that YOLOv8 achieves 50.2% mAP on the COCO dataset, an increase of approximately 4.5% over its predecessor YOLOv5, with an inference speed of up to 105FPS (GTX1060 platform). It supports N / S / M / L / X multi-scale model deployment, covering the full range of scenarios from edge devices to GPU servers. The framework provides a Python API and CLI toolchain, supports ONNX / TensorRT format export, and has been successfully applied to complex scenarios such as autonomous driving (vehicle detection accuracy reaches 98.7%), industrial quality inspection (0.01mm level defect recognition), and drone aerial photography. Through optimization methods such as knowledge distillation and channel pruning, the model parameters are reduced by 30% while maintaining an accuracy of over 95%, becoming a benchmark solution in the field of real-time target detection.

[0004] Academics currently primarily define small objects in terms of relative scale and absolute scale. For relative scale, common definitions include the ratio of the target bounding box area to the image area, the aspect ratio, and the pixel coverage ratio. For absolute scale, the commonly used definition comes from the MS COCO dataset, which defines small objects as those with a resolution less than 32 pixels × 32 pixels. Other datasets also have similar definitions. Existing algorithms primarily focus on the detection performance of large and medium-sized objects, with less optimization for small objects. Mainstream anchor-based detectors perform poorly for small objects and are prone to missed detections. Adjusting the anchor size also increases computational complexity. Summary of the Invention

[0005] This paper proposes an improved method for small target detection in remote sensing images based on YOLOv8. By introducing a Haar wavelet transform downsampling module (HWD), a wavelet transform feature enhancement module (WTFEM), and a minimum point distance-based bounding box regression loss function (MPDIoU), the method improves the detection accuracy and robustness of small targets. First, the downsampling module in YOLOv8 is replaced with a Haar wavelet transform downsampling module (HWD). HWD utilizes the Haar wavelet transform to reduce the spatial resolution of feature maps while preserving as much information as possible. This is crucial in remote sensing applications, as small targets can be obscured by background noise, and preserving spatial information is crucial for accurate detection.

[0006] Secondly, the Wavelet Transform Feature Enhancement Module (WTFEM) is innovatively introduced into the neck portion of the network, aiming to overcome the limitations of traditional feature fusion methods and significantly improve the model's ability to understand the semantics and retain details of multi-scale targets. WTFEM decomposes input features into low-frequency approximate components and high-frequency detail components through discrete wavelet transforms, extracts time-domain and frequency-domain features respectively using learnable asymmetric convolution kernels, and then reconstructs the enhanced multi-scale feature map through inverse transforms. While maintaining computational efficiency, this module enhances the ability to represent high-frequency information such as edge contours and texture details, making it particularly suitable for target detection tasks in remote sensing images with large target scales and complex backgrounds.

[0007] Finally, the original loss function was replaced with the bounding box regression loss function MPDIoU, based on minimum point distance. This loss function calculates the Euclidean distance of the nearest vertex between the predicted box and the true box, combining area overlap and center point distance optimization objectives to address the vanishing gradient problem of traditional IoU in scenarios where the bounding boxes do not overlap. MPDIoU not only quantifies the degree of spatial alignment but also dynamically adjusts regression weights through a minimum distance constraint, enhancing the ability to locate distant targets. It demonstrates significant advantages in the detection of densely packed small objects and objects with very low overlap. This improvement effectively alleviates boundary degradation, bringing the detection box closer to the actual target location, and comprehensively improving the model's positioning accuracy and generalization performance.

[0008] The specific implementation steps of the present invention are as follows:

[0009] Step 1: Standardize the input image (such as normalization, resizing, etc.) to meet the input requirements of the model.

[0010] Step 2: The input image is downsampled through the improved YOLOv8 model's backbone network to extract features at different scales. These downsampling modules have been replaced with Haar wavelet transform (HWD) downsampling modules. HWD utilizes the Haar wavelet transform to reduce the spatial resolution of the feature map while preserving as much information as possible. This is crucial in remote sensing applications, as small objects can be obscured by background noise, and preserving spatial information is crucial for accurate detection.

[0011] Step 3: After the image is downsampled using the Haar wavelet transform, feature maps at all levels are fed into the neck. In the neck, the wavelet feature enhancement module (WTFEM) fuses and enhances the multi-level features to improve the ability to extract small objects. The features processed by WTFEM are further sampled and fused in the neck of the original YOLOv8 to further extract feature information of small objects. The neck downsampling module is also replaced with Haar wavelet downsampling to improve the network's ability to capture spatial information.

[0012] Step 4: The processed feature map enters the head. This head performs object classification and bounding box regression prediction, outputting category confidence parameters, including the object's center coordinates, width, and height. Here, the network's original loss function is replaced with the bounding box regression loss function based on minimum point distance (MPDIoU).

[0013] Step 5: Apply the non-maximum suppression (NMS) algorithm to the output of the head part to remove redundant detection results and retain only the detection boxes with the highest confidence, ensuring the accuracy and reliability of the detection results. Finally, the processed object detection results, including the object category and bounding box, are output for subsequent applications. BRIEF DESCRIPTION OF THE DRAWINGS

[0014] Figure 1 Figure 2: Overall network structure diagram of the algorithm.

[0015] Figure 2 YOLOv8 structure diagram.

[0016] Figure 3 For: HWD structure diagram.

[0017] Figure 4 For: WTFEM structure diagram.

[0018] Figure 5 Figure 1: MPDIoU diagram. DETAILED DESCRIPTION

[0019] The present invention will be further described below with reference to the accompanying drawings.

[0020] The present invention designs a remote sensing image rotation small target detection network structure based on improved YOLOv8 (its structure is as follows Figure 1 As shown). We build on the foundation of YOLOv8 (its structure is as follows Figure 2 As shown), first, replace the backbone and neck downsampling modules in YOLOv8 with Haar wavelet downsampling modules HWD (its structure is as follows Figure 3 This sampling module uses the Haar wavelet transform to reduce the spatial resolution of the feature map while retaining as much information as possible, which is crucial in remote sensing applications because small targets may be obscured by background noise, and preserving spatial information is key to accurate detection.

[0021] Secondly, the wavelet feature enhancement module WTFEM (its structure is as follows) is integrated into the neck part of the network. Figure 4 This module strengthens the feature extraction capability and further improves the detection accuracy of small targets.

[0022] Finally, the original loss function was replaced with the bounding box regression loss function MPDIoU based on minimum point distance. This loss function calculates the Euclidean distance between the nearest vertex of the predicted box and the ground-truth box, combining the area overlap ratio and center point distance optimization objectives to solve the gradient vanishing problem of traditional IoU in scenarios where the bounding boxes do not overlap.

[0023] Specific methods

[0024] In this paper, an improved method for detecting rotating small targets based on YOLOv8 is proposed. By introducing the Haar wavelet transform downsampling module HWD, the wavelet transform feature enhancement module WTFEM and the bounding box regression loss function MPDIoU based on the minimum point distance, high-precision detection of small targets in remote sensing images is achieved.

[0025] Figure 1 This is the overall network structure diagram of the algorithm. First, the downsampling module in YOLOv8 is replaced with the Haar wavelet transform downsampling module HWD. HWD uses Haar wavelet transform to reduce the spatial resolution of the feature map while retaining as much information as possible. Secondly, the wavelet transform feature enhancement module WTFEM is innovatively introduced in the neck part of the network, aiming to break through the limitations of traditional feature fusion methods and significantly improve the model's semantic understanding and detail retention capabilities for multi-scale targets. Finally, the original loss function is replaced with the bounding box regression loss function MPDIoU based on minimum point distance. This loss function solves the problem of gradient vanishing of traditional IoU in scenarios where bounding boxes do not overlap by calculating the Euclidean distance of the nearest vertex between the predicted box and the true box, combining the area overlap rate and center point distance optimization objectives.

[0026] Figure 2 This is the YOLOv8 structure diagram. YOLOv8 is a new generation target detection framework launched by Ultralytics. It continues the core advantages of the YOLO series of single-stage detectors and achieves a significant breakthrough between speed and accuracy. Its architecture adopts a modular design. The backbone network (Backbone) is composed of an improved CSPDarknet53. The C2F module replaces the traditional C3 structure to enhance gradient flow and reduce parameter redundancy. At the same time, the SPPF layer is introduced to enhance multi-scale feature extraction capabilities. The detection head adopts a decoupled head structure to separate the classification and regression tasks, and for the first time introduces the Anchor-Free mechanism to directly predict the target center point, avoiding the complexity of anchor box size adaptation in traditional anchor-based methods.

[0027] Figure 3This is the HWD structure diagram. HWD mainly consists of two blocks: the lossless feature encoding block and the feature representation learning block. The lossless feature encoding block is responsible for converting features and reducing the spatial resolution. To achieve this goal, we utilize the Haar wavelet transform, which is a method that effectively reduces the resolution of feature maps while retaining all information. The representation learning block consists of a standard convolutional layer, batch normalization, and a ReLU activation layer. It is used to extract distinctive features. The lossless feature encoding block utilizes the Haar wavelet transform layer to effectively reduce the spatial resolution of the feature map while retaining all information. The Haar wavelet transform is a widely recognized compact, binary, orthogonal transform that is widely used in image coding, edge extraction, and binary logic design. The wavelet basis function and scaling function of the first-order one-dimensional Haar transform can be defined as follows:

[0028] φ1(x)=φ0(2x)+φ0(2x-1) (1)

[0029] ψ1(x)=φ0(2x)-φ0(2x-1) (2)

[0030] In this stage, the input signal or image is decomposed into two parts: low-frequency components and high-frequency components. The low-frequency components contain most of the information of the image, while the high-frequency components capture the details.

[0031] The feature representation learning block consists of a standard 1×1 convolutional layer, a batch normalization layer, and a ReLU activation function. In this block, standard convolution is used to adjust the number of channels in the feature map. This block has two main purposes: to adjust the number of channels in the feature map to align with subsequent layers; and to filter out as much redundant information as possible, enabling subsequent layers to more effectively learn representative features.

[0032] Figure 4 It is a WTFEM structure diagram. Due to the complexity of remote sensing images, the problem of false alarms with similar features is prone to occur in the task of detecting small rotating targets. At the same time, the feature extraction capability of existing methods is limited, and the extracted features contain less semantic information and a narrower receptive field, which makes it difficult to distinguish small rotating targets from the background. In order to solve this problem, the present invention adopts a multi-branch convolution structure to extract multiple discriminative semantic information from the perspective of improving feature richness; from the perspective of expanding the receptive field, the dilated convolution technology is applied to obtain richer local context information. Each branch performs a 1×1 convolution operation on the input feature map and preliminarily adjusts the number of channels for subsequent processing. The first branch is a residual structure, which is used to retain the key feature information of small targets. The other three branches perform cascaded standard convolution operations, and the wavelet convolution kernel sizes are 1×3, 3×1 and 3×3 respectively. Adding additional dilated convolution layers to the two middle branches enables the extracted feature maps to retain more contextual information, thereby more effectively identifying and distinguishing small targets. The calculation expression is as follows:

[0033]

[0034] Z=Concat(K1,K2,K3) (6)

[0035] Figure 5 This is a schematic diagram of MPDIoU. This loss function solves the problem of vanishing gradients in traditional IoU when bounding boxes do not overlap by calculating the Euclidean distance of the nearest vertex between the predicted box and the true box, combined with the area overlap rate and center point distance optimization objectives. MPDIoU not only quantifies the degree of spatial alignment, but also dynamically adjusts the regression weights through the minimum distance constraint, thereby enhancing the ability to locate distant targets, especially in the detection tasks of dense small targets and extremely low overlapping targets. This improvement effectively alleviates the boundary degradation phenomenon, makes the detection box closer to the actual position of the target, and comprehensively improves the positioning accuracy and generalization performance of the model. Its calculation formula is as follows:

[0036]

Claims

1. A rotating small target detection network based on improved YOLOv8 OBB, characterized by The following steps are involved: Step 1: The input image passes through the backbone network of the modified YOLOv8 OBB model. The downsampling module in the backbone has been replaced with a Haar wavelet downsampling module (HWD). The HWD downsampling module uses the Haar wavelet transform to reduce the spatial resolution of the feature map while preserving as much information as possible. This is crucial in remote sensing applications, as small targets can be obscured by background noise, and preserving spatial information is key to accurate detection. Step 2: After the image features are processed by the Haar wavelet transform downsampling module HWD in the backbone, the feature maps of the three scales enter the neck part. In the neck part, the wavelet transform feature enhancement module WTFEM is used to fuse and enhance the multi-level features to improve the feature extraction capability of small targets. The features processed by WTFEM are further sampled and fused through the neck part of the original YOLOv8 OBB to further extract the feature information of small targets. At the same time, the downsampling module of the neck part is also replaced with the Haar wavelet transform downsampling module HWD to retain as much spatial information as possible. In this network, we use the MPDIoU loss function to enhance the network's detection capability for multi-scale targets; Step 3: The processed feature map enters the head part. The head part performs object classification and bounding box regression prediction, and outputs the parameters of the category confidence, including the center coordinates, width, and height of the object.

2. The method according to claim 1, wherein: The downsampling module of the backbone in step 1 has been replaced by a Haar wavelet transform downsampling module HWD as a solution.

3. The method according to claim 1, wherein: In step 2, the wavelet transform feature enhancement module WTFEM is used to enhance the features of small targets.

4. The method according to claim 1, wherein: In step 2, we use the MPDIoU loss function to enhance the network's ability to detect multi-scale objects.

5. The method according to claim 1, wherein: In step 3, the head part outputs the category confidence and bounding box parameters, including the center coordinates, width, and height of the target.

Images