Millimeter wave image target segmentation method based on YOLOv8

Abstract

The invention discloses a millimeter wave image segmentation method based on YOLOv8. According to the method, an RFCACConv module is introduced according to the problem of parameter sharing of an original convolution kernel in an existing YOLOv8 structure, and spatial attention weight is added for the original convolution kernel by capturing spatial position information of a target so as to improve expression of feature capability. And by adjusting the target detection head, the information capture of the features of the small target is increased. In order to solve the problem that in an original sample selection strategy, a fixed hyper-parameter IoU threshold value is used for distributing positive and negative samples, and small target samples cause violent change of IoU loss, a label distribution strategy of a self-adaptive threshold value is adopted, the threshold value is dynamically adjusted according to target distribution characteristics of a data set, and the matching process of the positive and negative samples is optimized. According to the method, the contour information of the contraband in the millimeter wave image is utilized, so that working personnel can identify different types of contrabands more accurately.

Description

Technical Field

[0001] The present invention belongs to the technical field of millimeter wave image target detection, and in particular relates to a millimeter wave image target segmentation method based on YOLOv8. Background Art

[0002] Millimeter wave image detection systems have the ability to penetrate clothing and other surface objects, and can obtain information hidden under the surface of these objects. This technology is particularly useful in security inspections, and can detect potentially dangerous items, such as weapons and other illegal items, without violating privacy. Millimeter wave imaging has become an important solution for modern security inspections due to its non-contact, efficient and privacy-friendly characteristics. In the past few years, research on hidden objects in millimeter wave images has attracted widespread attention. Previous methods mainly rely on machine learning or mathematical models, but are limited by the constraints of small data sets, such as few locations and features, simple backgrounds, and limited object scales. The presence of noisy backgrounds in millimeter wave images, the high similarity between objects and backgrounds, the small size of most targets, and the impact of dynamic changes in the human body on imaging quality make it difficult for traditional methods to meet actual needs. In recent years, with the successful application of deep learning in the field of computer vision, deep learning-based methods have made significant progress in millimeter wave object detection. However, the current mainstream research focuses on bounding box-based detection methods. Although this method can effectively identify target objects, it often faces performance bottlenecks when dealing with small targets. Especially in millimeter wave images, targets are usually small and complex. In addition, the existing methods make insufficient use of the contour information of the target, which limits the accurate capture of the shape and edge of the object. In the millimeter wave image segmentation task, the application of deep learning methods is relatively insufficient. Therefore, the present invention proposes to use advanced deep learning technology to process millimeter wave images, especially to make breakthroughs in segmenting hidden objects. The use of segmentation technology can extract the contour information of the object, which is more helpful for staff to identify contraband such as guns and knives, thereby improving the visualization effect of the detection system. Summary of the invention

[0003] In view of the above problems, the present invention provides a millimeter wave segmentation method based on YOLOv8, which improves the visualization effect of the millimeter wave image detection system by utilizing the contour information of the target. First, the present invention introduces the receptive field attention convolution (RFCAConv) module to enhance the extraction of target contour information by combining spatial features. In view of the problem of IoU sensitivity of small targets, an adaptive threshold label assignment method is adopted to further ensure the model's accurate recognition of small target contours.

[0004] A millimeter wave image segmentation method based on YOLOv8 includes the following steps:

[0005] Step (1), collect millimeter wave image segmentation data sets, classify them according to contraband and add corresponding labels.

[0006] Step (2), preprocessing the millimeter wave image segmentation data set;

[0007] The preprocessing includes color enhancement, image translation and scaling, and mosaic enhancement and mixed sample enhancement are also used to increase the diversity of the data set.

[0008] Step (3), constructing a millimeter wave image segmentation model based on YOLOv8;

[0009] The YOLOv8-based millimeter wave image segmentation model uses a receptive field attention convolution (RFCAConv) module to replace the standard convolution module in the backbone network of YOLOv8. The RFCAConv module mainly includes a feature generation module, a feature rearrangement module, an adaptive pooling and splicing module, and an attention weight calculation module. At the same time, the feature pyramid network structure of PANet is adjusted in the YOLOv8 model.

[0010] Step (4), assigning positive and negative samples through an adaptive label assignment strategy;

[0011] Step (5): train a millimeter wave image segmentation model based on YOLOv8 to obtain test results.

[0012] Furthermore, the specific operation of the RFCAConv module is as follows:

[0013] The feature generation module uses a convolutional layer, batch normalization, and ReLU activation function to generate features to obtain F gen , using a group convolution with a stride of 1, each input channel generates K 2 output channels.

[0014] F gen =ReLU(BN(Conv2d(X)))

[0015] Where X represents the input feature map and the input channel is C×H×W.

[0016] The feature rearrangement module uses the view function to generate the feature map F gen Size CK 2 ×H×W are rearranged into C×KH×KW to obtain the rearranged feature map F rearrange In order to apply the attention mechanism.

[0017] Adaptive pooling and concatenation modules rearrange the feature map F rearrangeAdaptive average pooling is performed in height and width directions to capture long-distance information and obtain feature maps X of size C×1×KW and C×KH×1 h and X w .

[0018] X h =AdaptiveAvgPool2d(F rearrange )

[0019] X w =AdaptiveAvgPool2d(F rearrange )

[0020] The obtained feature maps are concatenated and merged, point convolution and batch normalization are applied, and activated by the activation function h_swish for subsequent convolution operations.

[0021] The attention weight calculation module first divides the feature map output by the adaptive pooling and splicing module into height and width parts. Then the feature map is convolved along the height and width directions to calculate the attention weights of the height and width, and normalized to the range of (0,1) by the sigmoid function to obtain the height and width weights a h and a w .

[0022] a h =σ(Conv2d(X h ))

[0023] a w =σ(Conv2d(X w ))

[0024] Finally, the height and width weights a h and a w Multiply to get the attention weight Ai:

[0025] Ai=a h ×a w

[0026] Multiply the attention weight Ai by the rearranged feature map F rearrange Finally, the output result F of the RFCAConv module is obtained through a k×k convolution operation with a step size of k. out .

[0027] F out =Conv2d(F rearrange ,Ai)

[0028] Furthermore, the feature pyramid network structure of PANet is adjusted. Specifically, a shallow high-resolution feature map is extracted from the first C2f module in Backbone to obtain high-resolution features. Then, the P3 layer feature map in the original Neck is upsampled to increase its resolution to the same level as the shallow high-resolution feature map. Figure 1 The upsampled feature map is fused with the shallow high-resolution feature map through the concat operation, and then further processed by the C2f and convolution layers, and finally passed to the Head to detect small targets. The shallow high-resolution P2 layer feature map is introduced into the fusion process, so that the P3, P4, and P5 layers combine the high-resolution information enhanced by the P2 layer during the concat operation.

[0029] Furthermore, the specific method of step (4) is as follows:

[0030] For each true box on the input image, we first select possible positive samples from each feature pyramid layer, and specifically select k candidate boxes closest to the center of the true box through the L2 distance. For an L-layer feature pyramid, each true box will have k×L candidate boxes. Next, we calculate the intersection-over-union (IoU) between these candidate boxes and the true box g, i.e., D g :

[0031]

[0032] Among them B candidate ∩B g represents the intersection area of ​​the candidate box and the real box g, B candidate ∪B g Represents the union area of ​​the candidate box and the real box g. Then determine the mean and standard deviation m of the candidate box g and v g .

[0033]

[0034]

[0035] Where D g,i Represents the IoU value of the i-th candidate box and the real box g. Based on these statistics, the IoU threshold of the real box g is set, that is,

[0036] t g =m g +v g

[0037] Finally, filter out the values ​​whose IoU is greater than or equal to the threshold t gThe candidate box of is used as the final positive sample, and the rest of the candidate boxes are used as negative samples. It should be noted that the center of the positive sample must be inside the true box. In addition, if a candidate box is assigned to multiple true boxes at the same time, the one with the largest IoU value is selected, and the rest will be used as negative samples.

[0038] Furthermore, the specific method of step (5) is as follows:

[0039] The preprocessed dataset is divided into a training set and a test set. After data enhancement in step (2), the YOLOv8-based millimeter wave image segmentation model constructed in step (3) is trained using the training set, and the positive and negative samples are matched using the label assignment method in step (4). After multiple rounds of forward propagation and back propagation iterative optimization of the loss function using the training set, the model is finally trained and evaluated on the test set.

[0040] Save the best model weights during the training process, use the best model weights to predict the test set to get the result image, and finally compare whether the recognition effect and detection accuracy meet the expected requirements.

[0041] The beneficial effects of the present invention are as follows:

[0042] Based on the parameter sharing problem of the original convolution kernel in the existing YOLOv8 structure, the present invention introduces the RFCAConv module, which adds spatial attention weights to the original convolution kernel by capturing the spatial position information of the target to enhance the expression of feature capabilities. And by adjusting the target detection head, the information capture of the features of small targets is increased.

[0043] In order to solve the problem that a fixed hyperparameter IoU threshold is used to allocate positive and negative samples in the original sample selection strategy, and small target samples cause drastic changes in IoU loss, the method of the present invention adopts an adaptive threshold label allocation strategy to dynamically adjust the threshold according to the target distribution characteristics of the data set, so as to optimize the matching process of positive and negative samples.

[0044] The present invention utilizes the contour information of contraband in the millimeter wave image, which is beneficial for staff to more accurately identify different types of contraband. BRIEF DESCRIPTION OF THE DRAWINGS

[0045] Figure 1 This is a flow chart of millimeter wave image segmentation based on the YOLOv8 model of the present invention.

[0046] Figure 2 Schematic diagram of the feature pyramid network structure in PANET.

[0047] Figure 3 A set of comparison pictures of the detection effect.

[0048] Figure 4 A set of comparison graphs of detection results for small objects. DETAILED DESCRIPTION

[0049] The technical solution of the present invention is further described below in conjunction with the accompanying drawings and embodiments.

[0050] A millimeter wave image segmentation method based on YOLOv8 includes the following steps:

[0051] Step (1), first, during the data set collection process, ensure that the data set can truly reflect the actual application scenario, the targets are randomly and evenly distributed in different positions in the image, and the gender ratio of the models is balanced. The data set is segmented by combining the pre-trained SAM (Segment Anything Model) model with manual annotation. The SAM model is first used for preliminary automatic segmentation, and then fine-tuned and proofread by manual. The data set contains a total of 4,000 images, of which 324 are background samples. A total of 5,500 targets are annotated in the data set, and these targets are divided into 6 categories, namely knives, explosives, liquids, guns, lighters and others. The data set is randomly divided into training and test sets in an 8:2 ratio.

[0052] Step (2), preprocessing the millimeter wave image segmentation data set;

[0053] The color effect of the image was changed by adjusting the hue of the image to 0.015, the saturation to 0.7, and the brightness to 0.4 to create more variations; secondly, new perspectives and size changes were generated by setting the image translation ratio to 0.1 and the scaling gain to 0.5. In addition, mosaic enhancement was applied, with a probability of 1.0 to partially mask the image area, specifically by randomly splicing 4 images into one image, so that parts of content from different samples are included in one image at the same time. Increasing the diversity of the data further expands the variability of the dataset.

[0054] Step (3), constructing a millimeter wave image segmentation model based on YOLOv8;

[0055] The YOLOv8-based millimeter wave image segmentation model uses the receptive field attention convolution (RFCAConv) module to replace the standard convolution module in the backbone network of YOLOv8, which hardly increases any computational cost and parameters, and can provide effective attention weights for the convolution kernel with spatial features, significantly improving network performance. Since the original convolution kernels share the same parameter K, the standard convolution operation cannot capture the information differences brought by different positions, which ultimately limits the performance of the network to a certain extent. The RFCAConv module mainly includes a feature generation module, a feature rearrangement module, an adaptive pooling and splicing module, and an attention weight calculation module.

[0056] like Figure 1 As shown, modify and configure the network structure of YOLOv8; reference the RFCAConv convolution module. The specific steps of RFCAConv are as follows:

[0057] The feature generation module uses a convolutional layer, batch normalization, and ReLU activation function to generate features to obtain F gen , here a group convolution with a step size of 1 is used, and each input channel generates K 2 output channels.

[0058] F gen =ReLU(BN(Conv2d(X)))

[0059] Where X represents the input feature map and the input channel is C×H×W.

[0060] The feature rearrangement module uses the view function to generate the feature map F gen Size CK 2 ×H×W are rearranged into features of the form C×KH×KW to obtain the rearranged feature map F rearrange In order to apply the attention mechanism.

[0061] Adaptive pooling and concatenation modules rearrange the feature map F rearrange Adaptive average pooling is performed in height and width directions to capture long-distance information and obtain feature maps X of size C×1×KW and C×KH×1 h and X w .

[0062] X h =AdaptiveAvgPool2d(F rearrange )

[0063] X w =AdaptiveAvgPool2d(F rearrange )

[0064] The obtained feature maps are concatenated and merged, point convolution and batch normalization are applied, and activated by the activation function h_swish for subsequent convolution operations.

[0065] The attention weight calculation module first divides the feature map output by the adaptive pooling and splicing module into height and width parts. Then the feature map is convolved along the height and width directions to calculate the attention weights of the height and width, and normalized to the range of (0,1) by the sigmoid function to obtain the height and width weights a h and a w .

[0066] a h =σ(Conv2d(X h ))

[0067] a w =σ(Conv2d(X w ))

[0068] Finally, the height and width weights a h and a w Multiply to get the attention weight Ai:

[0069] Ai=a h ×a w

[0070] The attention weight Ai is then multiplied by the rearranged feature map F rearrange , and finally the output result is obtained through a k×k convolution operation with a step size of k.

[0071] F out =Conv2d(F rearrange ,Ai)

[0072] like Figure 2 The traditional PANet feature pyramid network structure shown in the figure usually contains P3, P4, and P5 feature layers, which correspond to targets of different sizes: the P3 layer is responsible for detecting small targets, the P4 layer processes medium-sized targets, and the P5 layer is mainly used to detect large targets. However, due to the high proportion of small targets and even tiny targets in millimeter-wave images, relying solely on the P3 layer may not be enough to capture these fine feature information. Therefore, by introducing the P2 layer, higher-resolution features can be obtained to capture more subtle target details. Specifically, a shallow high-resolution feature map is extracted from the first C2f module in Backbone to obtain high-resolution features, and then the P3 layer feature map in the original Neck is upsampled to increase its resolution to the same level as the shallow high-resolution features. Figure 1The upsampled feature map is fused with the shallow high-resolution feature map through the concat operation, and then further processed by C2f and convolution layers, and finally passed to the Head to detect small targets. The shallow high-resolution P2 layer feature map is introduced into the fusion process, so that the P3, P4, and P5 layers combine the high-resolution information enhanced by the P2 layer during the concat operation. This allows each subsequent Head to benefit from more diverse features when processing targets of different scales, which helps to enhance the recognition ability of targets of different scales.

[0073] Step (4) allocates positive and negative samples through an adaptive label allocation strategy. The specific steps are as follows:

[0074] 4-1. Input the real bounding box set G, the number of feature pyramid layers L, and the i-th layer candidate box set E i , all candidate box sets E and parameter k.

[0075] 4-2. Initialize the positive sample set P and the negative sample set N.

[0076] 4-3. Perform the following operations for each ground-truth box g∈G:

[0077] Construct an empty set C g Used to store candidate positive samples associated with the true bounding box g.

[0078] For each level i in the range [1, L], use the L2 distance-based rule to select from the set E i Select k candidate boxes from the set S i , the centers of these candidate boxes are closest to the center of the ground-truth bounding box g.

[0079] The selected candidate box set S i Add to C g middle.

[0080] 4-4. Calculate the set C g The IoU with the real box g is obtained g :

[0081]

[0082] Among them B candidate ∩B g represents the intersection area of ​​the candidate box and the real box g, B candidate ∪B g Represents the union area of ​​the candidate box and the true box g.

[0083] 4-5. Calculate D g The average value m g and standard deviation v g .

[0084]

[0085]

[0086] Where D g,i Represents the IoU value between the i-th candidate box and the real box g.

[0087] 4-6. Calculate the IoU threshold t of the true bounding box g g :

[0088] t g =m g +v g

[0089] 4-7. For each candidate sample c∈C g Perform the following: If IoU(c,g)≥t g And the center of c is in g, add c to the positive sample set P.

[0090] 4-8. Set the unselected candidate box set as the negative sample set N=EP;

[0091] 4-9. Return the positive sample set P and the negative sample set N.

[0092] The average value m g and standard deviation v g The sum of is taken as the IoU threshold. The mean IoU value m of the object g Used to evaluate whether the preset candidate box matches the object. g indicates that the candidate box quality is good, so the IoU threshold should be set higher. On the contrary, low m g This means that the quality of most candidate boxes is poor, and the IoU threshold should be lowered. In addition, the IoU standard deviation v g It is used to measure which pyramid layers are more suitable for detecting the object. g Indicates that a specific pyramid layer is more suitable for the object, so by changing v g Add to m g To obtain a high threshold, select positive samples from this layer. Low v g This indicates that multiple pyramid levels are suitable for the object, so by changing v g Add to m g Get a low threshold and select positive samples from these layers. Using the mean m g and standard deviation v g The sum is used as the IoU threshold t g, can adaptively select a sufficient number of positive samples from the appropriate pyramid layer according to the statistical characteristics of the object. At the same time, it is required that the center of the positive sample must be inside the object. Candidate boxes outside the object will predict features outside the object, so they are not suitable as positive samples and should be removed during training.

[0093] Step (5): train a millimeter wave image segmentation model based on YOLOv8 to obtain test results.

[0094] The preprocessed dataset is divided into a training set and a test set. After data enhancement in step (2), the YOLOv8-based millimeter wave image segmentation model constructed in step (3) is trained using the training set, and the positive and negative samples are matched using the label assignment method in step (4). After multiple rounds of forward propagation and back propagation iterative optimization of the loss function using the training set, the model is finally trained and evaluated on the test set.

[0095] The trained and saved best.pt model weights are used to predict the test set using predict.py. In order to verify the effectiveness of the improvement, the performance of various mainstream image segmentation methods on the millimeter wave segmentation dataset is compared. The results are shown in the following table:

[0096]

[0097] Experimental results show that the proposed method improves the mAP@0.5 performance of the bounding box by 2.6% and the mAP@0.5 of the segmentation mask by 10.4% compared with the original YOLOv8 model on the millimeter wave segmentation dataset. Figure 3 and Figure 4 As shown in the figure, other models can identify and classify "relatively large objects with clear edges". However, for "small objects with low frequency", the detection results of different models are quite random. For fuzzy and small objects, different models often miss detection, misclassify, and segment edges that are quite different from the original image. In this case, the present invention shows relatively good performance.

[0098] It should be noted that the above embodiments can be freely combined as needed. The above description is only a detailed description of the preferred embodiments and principles of the present invention, but is not intended to limit the present invention. For those of ordinary skill in the art, according to the ideas provided by the present invention, there will be changes in the specific implementation methods, and these changes should also be regarded as the protection scope of the present invention.

Claims

1. A millimeter wave image segmentation method based on YOLOv8, characterized in that: The steps include: Step (1), collecting millimeter wave image segmentation data sets, and classifying them according to contraband and adding corresponding labels; Step (2), preprocessing the millimeter wave image segmentation data set; The preprocessing includes color enhancement, image translation and scaling, and mosaic enhancement and mixed sample enhancement are also used to increase the diversity of the data set; Step (3), constructing a millimeter wave image segmentation model based on YOLOv8; The YOLOv8-based millimeter wave image segmentation model uses a receptive field attention convolution (RFCAConv) module to replace the standard convolution module in the backbone network of YOLOv8. The RFCAConv module mainly includes a feature generation module, a feature rearrangement module, an adaptive pooling and splicing module, and an attention weight calculation module; at the same time, the feature pyramid network structure of PANet is adjusted in the YOLOv8 model; Step (4), assigning positive and negative samples through an adaptive label assignment strategy; Step (5): train a millimeter wave image segmentation model based on YOLOv8 to obtain test results.

2. The millimeter wave image segmentation method based on YOLOv8 according to claim 1, characterized in that: The specific operation of the RFCAConv module is as follows: The feature generation module uses a convolutional layer, batch normalization, and ReLU activation function to generate features to obtain F gen , using a group convolution with a stride of 1, each input channel generates K 2 output channels; F gen =ReLU(BN(Conv2d(X))) Where X represents the input feature map, and the input channel is C×H×W; The feature rearrangement module uses the view function to generate the feature map F gen Size CK 2 ×H×W are rearranged into C×KH×KW to obtain the rearranged feature map F rearrange In order to apply the attention mechanism; Adaptive pooling and concatenation modules rearrange the feature map F rearrange Adaptive average pooling is performed in height and width directions to capture long-distance information and obtain feature maps X of size C×1×KW and C×KH×1 h and X w ; X h =AdaptiveAvgPool2d(F rearrange ) X w =AdaptiveAvgPool2d(F rearrange ) The obtained feature maps are concatenated and merged, point convolution and batch normalization are applied, and activated by the activation function h_swish for subsequent convolution operations; The attention weight calculation module first divides the feature map output by the adaptive pooling and splicing module into height and width parts; then the feature map is convolved along the height and width directions to calculate the attention weights of the height and width, and normalized to the range of (0,1) by the sigmoid function to obtain the height and width weights a. h and a w ; a h =σ(Conv2d(Xh)) a w =σ(Conv2d(Xw)) Finally, the height and width weights a h and a w Multiply to get the attention weight Ai: You have=a h ×a w Multiply the attention weight Ai by the rearranged feature map F rearrange Finally, the output result F of the RFCAConv module is obtained through a k×k convolution operation with a step size of k. out ; F out =Conv2d(F rearrange ,Ai).

3. The millimeter wave image segmentation method based on YOLOv8 according to claim 1, characterized in that: The feature pyramid network structure of PANet is adjusted. Specifically, a shallow high-resolution feature map is extracted from the first C2f module in Backbone to obtain high-resolution features. Then, the P3 layer feature map in the original Neck is upsampled to increase its resolution to be consistent with the shallow high-resolution feature map, and the upsampled feature map is fused with the shallow high-resolution feature map through the concat operation. After further processing through the C2f and convolution layers, it is finally passed to the Head to detect small targets. The shallow high-resolution P2 layer feature map is introduced into the fusion process, so that the P3, P4, and P5 layers are combined with the high-resolution information enhanced by the P2 layer during the concat operation.

4. The millimeter wave image segmentation method based on YOLOv8 according to claim 1, characterized in that: The specific method of step (4) is as follows: For each true box on the input image, first select possible positive samples from each feature pyramid layer, and select k candidate boxes closest to the center of the true box through L2 distance; for L-layer feature pyramid, each true box will have k×L candidate boxes; next, calculate the intersection over union (IoU) between these candidate boxes and the true box g, i.e. D g : Among them B candidate ∩B g represents the intersection area of ​​the candidate box and the real box g, B candidate ∪B g Represents the union area of ​​the candidate box and the real box g; then determine the mean and standard deviation m of the candidate box g and v g ; Where D g,i Represents the IoU value of the i-th candidate box and the real box g; based on these statistics, the IoU threshold of the real box g is set, that is, t g =m g +v g Finally, filter out the values ​​whose IoU is greater than or equal to the threshold t g The candidate box of is taken as the final positive sample, and the remaining candidate boxes are taken as negative samples. It should be noted that the center of the positive sample must be inside the real box. In addition, if a candidate box is assigned to multiple ground-truth boxes at the same time, the one with the largest IoU value is selected and the rest are used as negative samples.

5. The millimeter wave image segmentation method based on YOLOv8 according to claim 1 or 4, characterized in that: The specific method of step (5) is as follows: The preprocessed data set is divided into a training set and a test set; after data enhancement in step (2), the YOLOv8-based millimeter wave image segmentation model constructed in step (3) is trained using the training set, and the positive and negative samples are matched using the label assignment method in step (4); after multiple rounds of forward propagation and back propagation iterative optimization of the loss function using the training set, the model is finally trained and evaluated on the test set; Save the best model weights during the training process, use the best model weights to predict the test set to get the result image, and finally compare whether the recognition effect and detection accuracy meet the expected requirements.

Images