Lightweight target detection method based on content awareness

Abstract

The invention relates to a lightweight target detection method based on content awareness, and belongs to the technical field of target detection. According to the method, data bottom layer information is increased by using an image adaptive gamma transformation algorithm, and the image readability is improved; on the basis of YOLOv5s, improved ShuffleNet V2 is used as a main network, and the quantity of model parameters is reduced; carrying out up-sampling by adopting a content awareness feature recombination operator, and enhancing an aggregation effect of context semantic information; a brand-new CBAMC3 module is designed for feature extraction so as to pay attention to feature information on a channel and a spatial position at the same time, and the feature expression accuracy is improved; and channel pruning is performed in the improved model, so that the parameter quantity is further reduced. According to the method, the real-time performance is improved while high precision is guaranteed, the model parameter quantity is reduced, and the generalization ability is high.

Description

Technical Field

[0001] The present invention belongs to the technical field of target detection and relates to a lightweight target detection method based on content perception. Background Art

[0002] Object detection is widely used in many fields such as intelligent manufacturing, autonomous driving, video surveillance, face recognition and security systems. In terms of quality control in the field of intelligent manufacturing, object detection is an important link to ensure product quality, appearance and safety. Timely and effective object detection is essential to maintaining product standards.

[0003] Traditional target detection technology mainly relies on image processing and machine vision technology, using cameras and sensors to monitor products in real time, combined with image processing algorithms to automatically identify and classify targets. Although traditional methods play an important role in target detection, they often face limitations when dealing with complex scenes and small target detection.

[0004] With the rapid development of artificial intelligence technology, deep learning has shown great potential in surface target detection of industrial products. By learning the defect characteristics in a large amount of data, deep learning has significantly improved the accuracy and efficiency of detection, making target detection more intelligent and automated. However, due to the different shapes and sizes of the detection targets, it is difficult to extract important features, resulting in low accuracy in small target detection. In addition, due to the relatively complex deep learning network model, large number of parameters, low lightweight, it needs to be deployed on equipment with strong computing power, which often leads to problems such as high cost, large equipment size, and difficulty in moving. Summary of the invention

[0005] In view of this, an object of the present invention is to provide a lightweight object detection method based on content perception.

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

[0007] A lightweight object detection method based on content perception, the method comprising the following steps:

[0008] S1. Obtain the target data set, perform data expansion, data enhancement, and data set division on the data set;

[0009] S2, using the improved ShuffleNet V2 as the backbone network of YOLOv5s to perform preliminary feature extraction on the dataset;

[0010] S3. In the neck network of YOLOv5s, the content-aware feature recombination operator CARAFE is used to upsample the output of the backbone network;

[0011] S4, use the CBAMC3 feature extraction module integrated with the CBAM attention mechanism to extract features from the upsampled data;

[0012] S5, the extracted features are input into several prediction heads of different scales in the prediction head network for target detection;

[0013] S6. Train and verify the YOLOv5s network model of the above process, and then prune the model to obtain a lightweight YOLOv5s model for real-time target detection.

[0014] Further, in step S1, the data set is first expanded by using geometric transformation, wherein the geometric transformation includes translation, flipping, and mirroring transformation;

[0015] Secondly, the image grayscale is increased by the adaptive gamma transform algorithm, where the adaptive gamma transform algorithm formula is:

[0016] V out (x,y)=cV in (x,y) γ(x,y)

[0017]

[0018] Among them, x and y represent the row coordinates and column coordinates of the image respectively, V in (x,y) and V out (x, y) represents the input grayscale value and output grayscale value at the position (x, y), γ(x, y) represents the gamma value dynamically adjusted according to the position (x, y), α and β are constants used to control the dynamic adjustment range of the gamma value. and Represents the minimum and maximum values ​​within a 5×5 neighborhood;

[0019] Finally, the dataset is divided into training set, validation set and test set according to the preset ratio.

[0020] Further, in step S2, the YOLOv5s network model includes an input end, a backbone network, a neck network and a prediction head. In the ShuffleNet V2 backbone network, a CBRM layer is first set, and then a combination of several spatial downsampling units Shuffle2, basic modules Shuffle1 and SE attention mechanism networks are connected, and finally a spatial pyramid pooling SPPF structure is set; several depthwise separable convolutions DWConv with a convolution kernel size of 5×5 are set in each module of ShuffleNet V2.

[0021] Among them, the basic module unit Shuffle1 is used to extract features, and the spatial downsampling unit Shuffle2 reduces the height and width of the input feature map by half and expands the number of channels.

[0022] Further, in step S3, the content-aware feature reconstruction operator CARAFE includes a kernel prediction module and a content-aware reconstruction module, wherein:

[0023] In the kernel prediction module, the following processes are included:

[0024] First, the input feature map of shape H×W×C is compressed to H×W×C using a 1×1 convolutional layer. m ;

[0025] Then it passes through the content encoder, using K encoder ×K encoder The convolution layer predicts the up-sampled kernel to generate a recombinant kernel with a size of K. up , the upsampling ratio is σ, and the shape is The feature map of , expands the channel dimension in the spatial dimension, and obtains a shape of The feature map of

[0026] Finally, the kernel normalizer is used to make the sum of the weight values ​​of the convolution kernel equal to 1;

[0027] In the content-aware reconstruction module, we transform the K up ×K up The region is convolved with the predicted upsampled kernel to get the output value.

[0028] Further, in step S4, the CBAMC3 module simultaneously derives the respective weight coefficients along the two independent dimensions of channel and space, and performs adaptive refinement. The specific process is as follows:

[0029] First, the input feature X is split into two branch features X 1 and X 2 , where X 2 Through the residual network BottleNeck with CBAM attention, and then through the Concat structure and X 2 Perform feature splicing, and the corresponding function expression is as follows:

[0030]

[0031] Z=y N +X 2 =F cbam (y N-1 )+y N-1 +X 2 ,y0 =X 1

[0032] in, represents feature separation, Z is the output feature of the CBAMC3 module, and y N represents the output features of two convolutional layers in the N residual network structure, F cbam Represents the output features after the CBAM attention mechanism;

[0033] The output feature of the CBAM module is expressed as:

[0034] Y=F cbam (X l )=M s (M c (X l )⊙X l )⊙(M c (X l )⊙X l )

[0035] Where Y represents the output feature after passing through the CBAM module, X l represents the input feature, M c represents the channel attention weight, M s represents the spatial attention weight and ⊙ represents the element-wise multiplication operation.

[0036] Further, in step S5, a plurality of prediction heads are set in the prediction head network of the YOLOv5s network, which are used to perform target detection on features of different scales, wherein the prediction head includes a first prediction head for small target detection, a second prediction head for medium target detection, and a third prediction head for large target detection.

[0037] Further, in step S6, the divided data set is used to train the YOLOv5s network. After the training is completed, the YOLOv5s network is lightweighted by the channel pruning algorithm, wherein the channel pruning algorithm uses the L1 regularization of the scaling factor γ of the BN layer to evaluate the importance of the input channel, which includes the following process:

[0038] First, sparse training is performed on the BN layer to make it sparse;

[0039] Then, the weight threshold is determined according to the number of BN layers that need to be retained, and the convolution channels whose absolute values ​​of scaling factors are less than the weight threshold are cropped;

[0040] Finally, the pruned model is fine-tuned and trained, and the optimal pruned model is obtained through continuous iteration.

[0041] Further, in step S6, the loss function of the channel pruning algorithm is as follows:

[0042]

[0043]

[0044] Among them, L task represents the loss term of the original task, L prun represents the pruning loss term; L pref represents the performance loss term; λ 1 and λ 2 Represent the weight coefficients of pruning loss and performance loss, W c is the weight of channel c;

[0045] The lightweight YOLOv5s network performs real-time detection of targets.

[0046] The beneficial effects of the present invention are:

[0047] 1. The present invention increases the grayscale of the image through an adaptive gamma transform algorithm, thereby improving the readability of the image and the classification accuracy after model training.

[0048] 2. The present invention uses the improved ShuffleNet V2 as the backbone network to reduce the number of model parameters and improve the detection speed; and uses CARAFE for upsampling to improve the receptive field while making it lightweight.

[0049] 3. The present invention proposes a CBAMC3 module for feature extraction, which enhances the weight sparsity in both channel and spatial dimensions, refines target features, and improves small target detection capabilities.

[0050] 4. The present invention further reduces the number of model parameters and improves real-time performance by performing channel pruning and fine-tuning on the improved model, making it easy to deploy on devices with weak computing power.

[0051] 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

[0052] 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:

[0053] Figure 1 A flowchart of a content-aware lightweight object detection method proposed by the present invention;

[0054] Figure 2 This is a schematic diagram of the improved YOLOv5s network model of the present invention;

[0055] Figure 3 This is a schematic diagram of the CBAMC3 module network used in the present invention;

[0056] Figure 4 Schematic diagram of the channel pruning process used in the present invention. DETAILED DESCRIPTION

[0057] 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.

[0058] Among them, the drawings are only used for illustrative explanations, and they only represent schematic diagrams rather than actual pictures, and should not be understood as limitations on the present invention. In order to better illustrate the embodiments of the present invention, some parts of the drawings may be omitted, enlarged or reduced, and do not represent the size of actual products. For those skilled in the art, it is understandable that some well-known structures and their descriptions in the drawings may be omitted.

[0059] The same or similar numbers in the drawings of the embodiments of the present invention correspond to the same or similar parts; in the description of the present invention, it should be understood that if the terms "upper", "lower", "left", "right", "front", "rear", etc. indicate the orientation or position relationship, they are based on the orientation or position relationship shown in the drawings, which is only for the convenience of describing the present invention and simplifying the description, rather than indicating or implying that the device or element referred to must have a specific orientation, be constructed and operate in a specific orientation. Therefore, the terms describing the position relationship in the drawings are only used for illustrative purposes and cannot be understood as limiting the present invention. For ordinary technicians in this field, the specific meanings of the above terms can be understood according to specific circumstances.

[0060] See also Figure 1 to Figure 4 , which is a lightweight object detection method based on content awareness.

[0061] Example

[0062] This embodiment provides an implementation process of a lightweight object detection method based on content perception, such as Figure 1 As shown, the method comprises the following steps:

[0063] Step S1: Obtain the target data set, expand the data set by geometric transformation, use the image adaptive gamma transformation algorithm to increase the underlying information of the data, and divide the data set into a training set, a validation set, and a test set;

[0064] Step S2: Use the improved ShuffleNet V2 as the backbone network of YOLOv5s, introduce the SE attention mechanism into ShuffleNet V2, and increase the convolution kernel size of the depthwise separable convolution;

[0065] Step S3: In the neck network of YOLOv5s, the content-aware feature recombination operator CARAFE is used for upsampling, and the upsampling process is dynamically adjusted according to the input feature content to improve the feature quality and model generalization ability during the upsampling process;

[0066] Step S4: Use the CBAMC3 feature extraction module integrated with the CBAM attention mechanism to extract features from the upsampled data;

[0067] Step S5: The extracted features are input into several prediction heads of different scales in the prediction head network for target detection;

[0068] Step S6: Use the data set for training and verification to obtain an improved model, and then perform channel pruning on the model to obtain a lightweight YOLOv5s model for real-time target detection.

[0069] In step S1, the printed circuit board defect dataset PCB_Dataset published by Peking University is used to expand the dataset using geometric transformations such as translation, flipping, and mirroring. The dataset is then divided into a training set, a validation set, and a test set in a ratio of 8:1:1. To solve the problem of similar image features between categories in the dataset, the image grayscale is increased by an adaptive gamma transform algorithm, thereby improving the readability of the image and the classification accuracy after model training. The formula of the adaptive gamma transform algorithm is as follows:

[0070] V out (x,y)=cV in (x,y) γ(x,y)

[0071]

[0072] Where x and y represent the row and column coordinates of the image, respectively, V in (x,y) and V out (x, y) represents the input grayscale value and output grayscale value at the position (x, y), γ(x, y) represents the gamma value dynamically adjusted according to the position (x, y), α and β are constants used to control the dynamic adjustment range of the gamma value. and Represents the minimum and maximum values ​​within a 5×5 neighborhood.

[0073] In step S2, see Figure 2 The YOLOv5s network model consists of four parts: input end, backbone network, neck network and prediction head. The improved ShuffleNet V2 is used as the backbone network. ShuffleNet V2 is mainly composed of Figure 2 The basic module Shuffle1 and the spatial downsampling Shuffle2 are composed of two units. The basic module unit only extracts features without changing the input feature map size and the number of channels, while the spatial downsampling unit halves the height and width of the input feature map and increases the number of channels to four times the original. The improved ShuffleNet V2 is used as the backbone network. The ShuffleNet V2 backbone network first includes a CBRM layer, then connects a combination of several spatial downsampling units Shuffle2, basic modules Shuffle1 and SE attention mechanism networks, and finally sets a spatial pyramid pooling (SPPF) structure. Among them, the CBRM layer, spatial downsampling unit Shuffle2, and basic module Shuffle1 can improve the model's detection ability for targets of different sizes; the SE attention mechanism that automatically learns the importance of feature channels after the basic module and spatial downsampling can improve the model's accuracy and generalization ability;

[0074] In addition, the convolution kernel size of the depthwise separable convolution (DWConv) in each module of ShuffleNet V2 is increased from 3×3 to 5×5 to compensate for the feature loss caused by DWConv and retain more features.

[0075] In step S3, see Figure 2 , upsampling is performed using a content-aware feature reconstruction operator, which mainly consists of two key components: a kernel prediction module and a content-aware reconstruction module;

[0076] The process of the kernel prediction module is as follows: first, the input feature map of shape H×W×C is compressed into H×W×C by using a 1×1 convolutional layer through the channel compressor. m , and then pass through the content encoder, using K encoder ×K encoder The convolution layer predicts the up-sampled kernel to generate a recombinant kernel with a size of K. up , the upsampling ratio is σ, and the shape is The feature map of , expands the channel dimension in the spatial dimension, and obtains a shape of The feature map is then passed through a kernel normalizer to make the sum of the weight values ​​of the convolution kernel equal to 1.

[0077] The content-aware module is: by taking the K centers at each position in the input feature map up ×K up The neighborhood is convolved with the predicted upsampled kernel to get the output value.

[0078] The corresponding function expression is as follows:

[0079] W i ′=ψ(N(X i ,K encoder ))

[0080] X i ′=φ(N(X i ,K up ),W i′ )

[0081] Among them, ψ and φ represent the kernel prediction module and content-aware reconstruction module respectively, and X i is the input feature, N(X i ,K) indicates that X i is the K×K sub-neighborhood centered on W i′ is the output feature of the kernel prediction module, X i ′ is the output feature of the entire module.

[0082] In step S4, see Figure 3 , the CBAM attention mechanism is combined with the residual network of the C3 module to form a new CBAMC3 module, which can derive the respective weight coefficients along the two independent dimensions of channel and space at the same time, and perform adaptive refinement. This module first splits the input feature X into two branch features X 1 and X 2 , where X 2 Through the residual network BottleNeck with CBAM attention, and then through the Concat structure and X 2 Perform feature splicing, and the corresponding function expression is as follows:

[0083]

[0084] Z=y N +X 2 =F cbam (y N-1 )+y N-1 +X 2 ,y 0 =X 1

[0085] in represents feature separation, Z is the output feature of the CBAMC3 module, and y Nrepresents the output features of two convolutional layers in the N residual network structure, F cbam Represents the output features after the CBAM attention mechanism.

[0086] The CBAM module is a simple and effective feedforward convolutional neural network that can derive the weight coefficients of the two independent dimensions of channel and space at the same time, and then adaptively refine the input features to improve the feature expression and generalization capabilities of the neural network. Its function can be expressed as:

[0087] Y=F cbam (X l )=M s (M c (X l )⊙X l )⊙(M c (X l )⊙X l )

[0088] Where Y represents the output feature after passing through the CBAM module, X l represents the input feature, M c represents the channel attention weight, M s represents the spatial attention weight and ⊙ represents the element-wise multiplication operation.

[0089] After the above algorithm is improved, the data set described in step S1 is used for training, verification and testing, and finally an improved algorithm model is obtained.

[0090] In step S5, this embodiment sets three prediction heads in the prediction head network, namely prediction head 1, prediction head 2 and prediction head 3, wherein prediction head 1 performs small target detection with a scale of 80*80, prediction head 2 performs medium target detection with a scale of 40*40, and prediction head 3 performs large target detection with a scale of 20*20.

[0091] In step S6, refer to Figure 4 , prune the channels with weights of 0.006 and 0.008 in the figure, retain important channels, and make the model lightweight. The channel pruning algorithm uses the L1 regularization of the scaling factor γ of the BN layer to evaluate the importance of the input channel. First, the BN layer is sparsely trained to make the BN layer sparse, and then the weight threshold is determined according to the number of BN layers that need to be retained. The convolution channels whose absolute value of the scaling factor is less than the weight threshold are pruned, and finally the pruned model is fine-tuned to compensate for the impact of pruning on the accuracy of the detection model, and it is continuously iterated to obtain the best pruned model. The loss function of channel pruning is as follows:

[0092]

[0093] Where L task represents the loss term of the original task, L prun represents the pruning loss term, which is used to introduce sparsity and make the weights of unimportant channels close to zero. pref represents the performance loss term, which is used to constrain the computational efficiency of the model and ensure that the model can run efficiently after pruning. 1 and λ 2 Represent the weight coefficients of pruning loss and performance loss, W c is the weight of channel c.

[0094] 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.

Claims

1. A content-aware lightweight object detection method, characterized in that: The method comprises the following steps: S1. Obtain the target data set, perform data expansion, data enhancement, and data set division on the data set; S2, using the improved ShuffleNet V2 as the backbone network of YOLOv5s to perform preliminary feature extraction on the dataset; S3. In the neck network of YOLOv5s, the content-aware feature recombination operator CARAFE is used to upsample the output of the backbone network; S4, use the CBAMC3 feature extraction module integrated with the CBAM attention mechanism to extract features from the upsampled data; S5, the extracted features are input into several prediction heads of different scales in the prediction head network for target detection; S6. Train and verify the YOLOv5s network model of the above process, and then prune the model to obtain a lightweight YOLOv5s model for real-time target detection.

2. The content-aware lightweight target detection method according to claim 1, characterized in that: In step S1, a data set is first expanded by using geometric transformation, wherein the geometric transformation includes translation, flipping, and mirroring transformation; Secondly, the image grayscale is increased by the adaptive gamma transform algorithm, where the adaptive gamma transform algorithm formula is: V out (x,y)=cV in (x,y) γ(x,y) Among them, x and y represent the row coordinates and column coordinates of the image respectively, V in (x,y) and V out (x, y) represents the input grayscale value and output grayscale value at the position (x, y), γ(x, y) represents the gamma value dynamically adjusted according to the position (x, y), α and β are constants used to control the dynamic adjustment range of the gamma value. and Represents the minimum and maximum values ​​within a 5×5 neighborhood; Finally, the dataset is divided into training set, validation set and test set according to the preset ratio.

3. The method for lightweight object detection based on content perception according to claim 2, characterized in that: In step S2, the YOLOv5s network model includes an input end, a backbone network, a neck network and a prediction head. In the ShuffleNetV2 backbone network, a CBRM layer is first set, and then a combination of several spatial downsampling units Shuffle2, basic modules Shuffle1 and SE attention mechanism networks are connected, and finally a spatial pyramid pooling SPPF structure is set; several depth-separable convolutions DWConv with a convolution kernel size of 5×5 are set in each module of ShuffleNet V2. Among them, the basic module unit Shuffle1 is used to extract features, and the spatial downsampling unit Shuffle2 reduces the height and width of the input feature map by half and expands the number of channels.

4. The content-aware lightweight target detection method according to claim 3, characterized in that: In step S3, the content-aware feature reconstruction operator CARAFE includes a kernel prediction module and a content-aware reconstruction module, wherein: In the kernel prediction module, the following processes are included: First, the input feature map of shape H×W×C is compressed to H×W×C using a 1×1 convolutional layer. m ; Then it passes through the content encoder, using K encoder ×K encoder The convolution layer predicts the up-sampled kernel to generate a recombinant kernel with a size of K. up , the upsampling ratio is σ, and the shape is The feature map of , expands the channel dimension in the spatial dimension, and the shape is The feature map of Finally, the kernel normalizer is used to make the sum of the weight values ​​of the convolution kernel equal to 1; In the content-aware reconstruction module, we transform the K up ×K up The region is convolved with the predicted upsampled kernel to get the output value.

5. The content-aware lightweight object detection method according to claim 4, characterized in that: In step S4, the CBAMC3 module simultaneously derives the respective weight coefficients along the two independent dimensions of channel and space for adaptive refinement. The specific process is as follows: First, the input feature X is split into two branch features X1 and X2, where X2 passes through the residual network BottleNeck with CBAM attention, and then the Concat structure is used to concatenate features with X2. The corresponding function expression is as follows: Z=y N +X2=F cbam (y N-1 )+y N-1 +X2,y0=X1 in, represents feature separation, Z is the output feature of the CBAMC3 module, and y N represents the output features of two convolutional layers in the N residual network structure, F cbam Represents the output features after the CBAM attention mechanism; The output feature of the CBAM module is expressed as: Y=F cbam (X l )=M s (M c (X l )⊙X l )⊙(M c (X l )⊙X l ) Where Y represents the output feature after passing through the CBAM module, X l represents the input feature, M c represents the channel attention weight, M s represents the spatial attention weight and ⊙ represents the element-wise multiplication operation.

6. The content-aware lightweight object detection method according to claim 5, characterized in that: In step S5, a plurality of prediction heads are set in the prediction head network of the YOLOv5s network, which are used to perform target detection on features of different scales, wherein the prediction heads include a first prediction head for small target detection, a second prediction head for medium target detection, and a third prediction head for large target detection.

7. The content-aware lightweight object detection method according to claim 6, characterized in that: In step S6, the divided data set is used to train the YOLOv5s network. After the training is completed, the YOLOv5s network is lightweighted by the channel pruning algorithm, where the channel pruning algorithm uses the L1 regularization of the scaling factor γ of the BN layer to evaluate the importance of the input channel, which includes the following process: First, sparse training is performed on the BN layer to make it sparse; Then, the weight threshold is determined according to the number of BN layers that need to be retained, and the convolution channels whose absolute values ​​of scaling factors are less than the weight threshold are cropped; Finally, the pruned model is fine-tuned and trained, and the optimal pruned model is obtained through continuous iteration.

8. The content-aware lightweight object detection method according to claim 7, characterized in that: In step S6, the loss function of the channel pruning algorithm is as follows: Among them, L task represents the loss term of the original task, L prun represents the pruning loss term; L pref represents the performance loss term; λ1 and λ2 represent the weight coefficients of pruning loss and performance loss respectively, W c is the weight of channel c; The lightweight YOLOv5s network performs real-time detection of targets.

Images