A lightweight object detection method based on content awareness

Abstract

This invention relates to a content-aware, lightweight object detection method, belonging to the field of object detection technology. This method enhances image readability by using an image adaptive gamma transform algorithm to increase underlying data information; it reduces the number of model parameters by using an improved ShuffleNet V2 backbone network based on YOLOv5s; it employs a content-aware feature recombination operator for upsampling to enhance the aggregation effect of contextual semantic information; a novel CBAMC3 module is designed for feature extraction to simultaneously focus on channel and spatial location feature information, improving feature representation accuracy; and channel pruning is performed in the improved model to further reduce the number of parameters. This method improves real-time performance and reduces the number of model parameters while maintaining high accuracy, and has strong generalization ability.

Description

Technical Field

[0001] This invention belongs to the field of object detection technology and relates to a content-aware, lightweight object detection method. Background Technology

[0002] Target inspection is widely used in various fields such as smart manufacturing, autonomous driving, video surveillance, facial recognition, and security systems. In the field of smart manufacturing, target inspection is a crucial step in quality control, ensuring product quality, appearance, and safety. Timely and effective target inspection is essential for maintaining product standards.

[0003] Traditional target detection technologies primarily rely on image processing and machine vision techniques, utilizing cameras and sensors to monitor products in real time and combining image processing algorithms to automatically identify and classify targets. While traditional methods have played a significant role in target detection, they often face limitations when handling complex scenes and detecting small targets.

[0004] With the rapid development of artificial intelligence technology, deep learning has shown great potential in the surface target inspection of industrial products. By learning defect features from large amounts of data, deep learning has significantly improved the accuracy and efficiency of detection, making target inspection more intelligent and automated. However, due to the different shapes and sizes of the targets to be detected, it is difficult to extract important features, resulting in low accuracy in the detection of small targets. Furthermore, because deep learning network models are relatively complex, have a large number of parameters, and are not lightweight, they need to be deployed on devices with strong computing power, which often leads to problems such as high cost, large device size, and difficulty in moving the equipment. Summary of the Invention

[0005] In view of this, the purpose of this invention is to provide a content-aware, lightweight object detection method.

[0006] To achieve the above objectives, the present invention provides the following technical solution:

[0007] A content-aware, lightweight object detection method, comprising the following steps:

[0008] S1. Obtain the target dataset, and perform data augmentation, data enhancement, and dataset partitioning on the dataset;

[0009] S2. An improved ShuffleNet V2 is used 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. The CBAMC3 feature extraction module, which incorporates the CBAM attention mechanism, is used 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 validate the YOLOv5s network model from the above process, and then perform channel pruning on the model to obtain a lightweight YOLOv5s model for real-time object detection.

[0014] Furthermore, in step S1, the dataset is first expanded using geometric transformations, which include translation, flipping, and mirroring.

[0015] Secondly, the image grayscale is increased using an adaptive gamma transform algorithm, the formula for which is:

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

[0017]

[0018] Where x and y represent the row and column coordinates of the image, respectively, and V in (x,y) and V out (x,y) represent the input and output grayscale values ​​at position (x,y), respectively; γ(x,y) represents the dynamically adjusted gamma value based on different positions (x,y); and α 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, validation, and test sets according to a preset ratio.

[0020] Furthermore, in step S2, the YOLOv5s network model includes an input layer, a backbone network, a neck network, and a prediction head. In the ShuffleNet V2 backbone network, a CBRM layer is first set up, followed by a combination of several spatial downsampling units (Shuffle2), the basic module (Shuffle1), and the SE attention mechanism network. Finally, a spatial pyramid pooling (SPPF) structure is set up. In each module of ShuffleNet V2, several depthwise separable convolutional layers (DWConv) with a kernel size of 5×5 are set up.

[0021] The basic module unit Shuffle1 is used to extract features, while the spatial downsampling unit Shuffle2 halves the height and width of the input feature map and increases the number of channels.

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

[0023] The kernel prediction module includes the following processes:

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

[0025] Then it goes through the content encoder, using K... encoder ×K encoder The convolutional layer predicts the upsampling kernel to generate a reconstructed kernel, and the size of the reconstructed kernel is K. up With an upsampling ratio of σ, the resulting shape is... The feature map is expanded in the spatial dimension to obtain a shape of... Feature map;

[0026] Finally, a kernel normalizer is used to make the sum of the weights of the convolution kernel equal to 1.

[0027] In the content-aware reconstruction module, K is centered at each location in the input feature map. up ×K up The region is convolved with the predicted upsampling kernel to obtain the output value.

[0028] Furthermore, in step S4, the CBAMC3 module simultaneously derives its respective weight coefficients along both the channel and spatial independent dimensions, performing adaptive refinement. The specific process is as follows:

[0029] First, the input feature X is split into two branch features X1 and X2. X2 is passed through a residual network BottleNeck with CBAM attention. Then, the features are concatenated with X2 through a Concat structure. The corresponding function expression is as follows:

[0030]

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

[0032] in, Indicates feature separation, Z is the output feature of the CBAMC3 module, y NF represents the output features of two convolutional layers in an N residual network structure. cbam This represents the output features after passing through the CBAM attention mechanism;

[0033] The output characteristics of the CBAM module are represented as follows:

[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 M represents the input features. c M represents the channel attention weight. s ⊙ represents spatial attention weights, and ⊙ represents element-wise multiplication operations.

[0036] Furthermore, in step S5, the YOLOv5s network has several prediction heads set in its prediction head network, which are used to perform target detection on features of different scales. 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.

[0037] Further, in step S6, the YOLOv5s network is trained using the partitioned dataset. After training, the YOLOv5s network is lightweighted using a channel pruning algorithm. The channel pruning algorithm uses L1 regularization of the scaling factor γ of the BN layer to evaluate the importance of the input channels, and includes the following process:

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

[0039] Then, determine the weight threshold based on the number of BN layers to retain, and prune convolutional channels whose absolute scaling factor is less than the weight threshold.

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

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

[0042]

[0043]

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

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

[0046] The beneficial effects of this invention are as follows:

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

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

[0049] 3. This invention proposes the CBAMC3 module for feature extraction, which enhances the weight sparsity in both channel and spatial dimensions, refines target features, and improves the detection capability of small targets.

[0050] 4. This 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 easier to deploy on devices with limited computing power.

[0051] Other advantages, objects, and features of the present invention will be described in part in the following description and, in part, will be apparent to those skilled in the art upon examination of the following description or may be learned from practice of the present invention. The objects and other advantages of the present invention may be realized and obtained through the following description. Attached Figure Description

[0052] To make the objectives, technical solutions, and advantages of the present invention clearer, the preferred embodiments of the present invention will be described in detail below with reference to the accompanying drawings, wherein:

[0053] Figure 1 This is a flowchart of a content-aware lightweight target detection method proposed in this invention;

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

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

[0056] Figure 4This is a schematic diagram of the channel pruning process used in this invention. Detailed Implementation

[0057] The following specific examples illustrate the implementation of the present invention. Those skilled in the art can easily understand other advantages and effects of the present invention from the content disclosed in this specification. The present invention can also be implemented or applied through other different specific embodiments, and various details in this specification can be modified or changed 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 are only schematic representations of the basic concept of the present invention. Unless otherwise specified, the following embodiments and features can be combined with each other.

[0058] The accompanying drawings are for illustrative purposes only and are schematic diagrams, not actual pictures. They should not be construed as limiting the invention. To better illustrate the embodiments of the invention, some parts in the drawings may be omitted, enlarged, or reduced, and do not represent the actual product dimensions. It is understandable to those skilled in the art that some well-known structures and their descriptions may be omitted in the drawings.

[0059] In the accompanying drawings of the embodiments of the present invention, the same or similar reference numerals correspond to the same or similar components. In the description of the present invention, it should be understood that if terms such as "upper," "lower," "left," "right," "front," and "rear" indicate the orientation or positional relationship based on the orientation or positional relationship shown in the drawings, they are only for the convenience of describing the present invention and simplifying the description, and do not indicate or imply that the device or element referred to must have a specific orientation, or be constructed and operated in a specific orientation. Therefore, the terms used to describe positional relationships in the drawings are only for illustrative purposes and should not be construed as limiting the present invention. For those skilled in the art, the specific meaning of the above terms can be understood according to the specific circumstances.

[0060] Please see Figures 1-4 This is a lightweight object detection method based on content awareness.

[0061] Example

[0062] This embodiment presents the implementation process of a content-aware lightweight object detection method, such as... Figure 1 As shown, the method includes the following steps:

[0063] Step S1: Obtain the target dataset, expand the dataset using geometric transformation, add underlying information to the data using the image adaptive gamma transform algorithm, and divide the dataset into training set, validation set and test set;

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

[0065] Step S3: In the neck network of YOLOv5s, the content-aware feature recombination operator CARAFE is used for upsampling. 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, which incorporates the CBAM attention mechanism, to extract features from the upsampled data;

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

[0068] Step S6: Use the dataset to train and validate the improved model, then perform channel pruning on the model to obtain a lightweight YOLOv5s model for real-time object detection.

[0069] In step S1, the publicly available printed circuit board defect dataset PCB_Dataset from Peking University is used. Geometric transformations such as translation, flipping, and mirroring are applied to expand the dataset. Then, the dataset is divided into training, validation, and test sets in an 8:1:1 ratio. To address the issue of image feature similarity between categories in the dataset, an adaptive gamma transform algorithm is used to increase image grayscale, thereby improving image readability and the classification accuracy after model training. The formula for 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) represent the input and output grayscale values ​​at position (x,y), respectively; γ(x,y) represents the dynamically adjusted gamma value based on different positions (x,y); and α and β are constants used to control the dynamic adjustment range of the gamma value. and This represents the minimum and maximum values ​​within a 5×5 neighborhood.

[0073] In step S2, please refer to Figure 2The YOLOv5s network model consists of four parts: the input network, the backbone network, the neck network, and the prediction head. It uses an improved ShuffleNet V2 as the backbone network, where ShuffleNet V2 is primarily composed of… Figure 2 The model consists of two units: the basic module Shuffle1 and the spatial downsampling unit Shuffle2. The basic module extracts features without changing the size or number of channels of the input feature map, 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. An improved ShuffleNet V2 is used as the backbone network. The ShuffleNet V2 backbone network first includes a CBRM layer, followed by combinations of several spatial downsampling units (Shuffle2), the basic module Shuffle1, and an SE attention mechanism network. Finally, a Spatial Pyramid Pooling (SPPF) structure is set up. The CBRM layer, the spatial downsampling unit Shuffle2, and the basic module Shuffle1 improve the model's ability to detect targets of different sizes. The SE attention mechanism, which automatically learns the importance of feature channels, introduced after the basic module and spatial downsampling, improves the model's accuracy and generalization ability.

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

[0075] In step S3, please refer to Figure 2 Upsampling is performed using a content-aware feature recombination operator, which mainly consists of two key components: a kernel prediction module and a content-aware recombination module.

[0076] The kernel prediction module works as follows: First, it passes through a channel compressor, and then uses a 1×1 convolutional layer to compress the input feature map of shape H×W×C into H×W×C. m Then it goes through the content encoder, using K encoder ×K encoder The convolutional layer predicts the upsampling kernel to generate a reconstructed kernel, and the size of the reconstructed kernel is K. up With an upsampling ratio of σ, the resulting shape is... The feature map is expanded in the spatial dimension to obtain a shape of... The feature map is then processed by a kernel normalizer to make the sum of the weights of the convolution kernel equal to 1.

[0077] The content-aware module is: by using K, centered at each location, in the input feature map. up ×K up The output value is obtained by convolving the neighborhood with the predicted upsampling kernel.

[0078] Its 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] Where ψ and φ represent the kernel prediction module and the content-aware reconstruction module, respectively, and X i For the input features, N(X) i K) represents X i The K×K sub-neighborhood centered on W i′ X is the output feature after passing through the kernel prediction module. i ′ represents the output feature of the entire module.

[0082] In step S4, please refer to Figure 3 By combining the CBAM attention mechanism into the residual network of the C3 module, a novel CBAMC3 module is formed. This module can simultaneously derive the weight coefficients along two independent dimensions, channel and space, for adaptive refinement. The module first splits the input feature X into two branch features X1 and X2. X2 is passed through the BottleNeck residual network with CBAM attention, and then concatenated with X2 through a Concat structure. The corresponding function expression is as follows:

[0083]

[0084] Z = y N +X2=F cbam (y N-1 )+y N-1 +X2,y0=X1

[0085] in Indicates feature separation, Z is the output feature of the CBAMC3 module, y N F represents the output features of two convolutional layers in an N residual network structure. cbam This represents the output features after passing through the CBAM attention mechanism.

[0086] The CBAM module is a simple and effective feedforward convolutional neural network that can simultaneously derive weight coefficients along two independent dimensions—channel and spatial—and then adaptively refine the input features, improving the feature representation 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 M represents the input features. c M represents the channel attention weight. s ⊙ represents spatial attention weights, and ⊙ represents element-wise multiplication operations.

[0089] After the algorithm was improved, the dataset described in step S1 was used for training, validation and testing to finally obtain the improved algorithm model.

[0090] In step S5, this embodiment sets up three prediction heads in the prediction head network, namely prediction head 1, prediction head 2 and prediction head 3. 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, please refer to Figure 4 The channel pruning algorithm prunes channels with weights of 0.006 and 0.008 in the graph, retaining important channels and making the model lighter. The algorithm uses L1 regularization of the scaling factor γ of the Batch Normalization (BN) layer to evaluate the importance of input channels. First, the BN layer is sparsely trained. Then, a weight threshold is determined based on the number of BN layers to retain. Convolutional channels with a scaling factor absolute value less than this threshold are pruned. Finally, the pruned model is fine-tuned to compensate for the decrease in detection model accuracy caused by pruning. The process is iterated until the optimal pruned model is obtained. The loss function for channel pruning is shown below:

[0092]

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

[0094] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention and are not intended to limit it. Although the present invention has been described in detail with reference to preferred embodiments, those skilled in the art should understand that modifications or equivalent substitutions can be made to the technical solutions of the present invention without departing from the spirit and scope of the present invention, and all such modifications or substitutions should be covered within 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 dataset, and perform data augmentation, data enhancement, and dataset partitioning on the dataset; S2. An improved ShuffleNet V2 is used 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. The CBAMC3 feature extraction module, which incorporates the CBAM attention mechanism, is used 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 validate the YOLOv5s network model from the above process, and then perform channel pruning on the model to obtain a lightweight YOLOv5s model for real-time object detection. In step S1, the dataset is first expanded using geometric transformations, which include translation, flipping, and mirroring. Secondly, the image grayscale is increased using an adaptive gamma transform algorithm, the formula for which is: V out (x,y)=c0V in (x,y) γ(x,y) Where x and y represent the row and column coordinates of the image, respectively, and V in (x,y) and V out (x,y) represent the input and output grayscale values ​​at position (x,y), respectively; γ(x,y) represents the dynamically adjusted gamma value based on different positions (x,y); and α 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; In step S4, the CBAMC3 module simultaneously derives the respective weight coefficients along two independent dimensions, channel and space, and performs adaptive refinement. The specific process is as follows: First, the input feature X is split into two branch features X1 and X2. X2 is passed through a residual network BottleNeck with CBAM attention. Then, the features are concatenated with X2 through a Concat structure. The corresponding function expression is as follows: Z=y N +X2=F cbam (y N-1 )+y N-1 +X2,y0=X1 in, Indicates feature separation, Z is the output feature of the CBAMC3 module, y N F represents the output features of two convolutional layers in an N residual network structure. cbam This represents the output features after passing through the CBAM attention mechanism; The output characteristics of the CBAM module are represented as follows: 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 M represents the input features. c M represents the channel attention weight. s represents spatial attention weights, and ⊙ represents element-wise multiplication operations; In step S6, the loss function of the channel pruning algorithm is as follows: Among them, L task L represents the loss term of the original task. prun L represents the pruning loss term; pref W represents the performance loss term; λ1 and λ2 represent the weighting coefficients of pruning loss and performance loss, respectively. c It is the weight of channel c; The lightweight YOLOv5s network performs real-time target detection.

2. The content-aware lightweight target detection method according to claim 1, characterized in that: In step S2, the YOLOv5s network model includes an input layer, a backbone network, a neck network, and a prediction head. In the ShuffleNetV2 backbone network, a CBRM layer is first set up, followed by a combination of several spatial downsampling units (Shuffle2), the basic module (Shuffle1), and an SE attention mechanism network. Finally, a spatial pyramid pooling (SPPF) structure is set up. Within each module of ShuffleNet V2, several depthwise separable convolutional layers (DWConv) with a kernel size of 5×5 are set up. The basic module unit Shuffle1 is used to extract features, while the spatial downsampling unit Shuffle2 halves the height and width of the input feature map and increases the number of channels.

3. The content-aware lightweight target detection method according to claim 2, 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, The kernel prediction module includes the following processes: First, the input feature map, initially of shape H×W×C, is compressed into H×W×C using a channel compressor and a 1×1 convolutional layer. m ; Then it goes through the content encoder, using K... encoder ×K encoder The convolutional layer predicts the upsampling kernel to generate a reconstructed kernel, and the size of the reconstructed kernel is K. up With an upsampling ratio of σ, the resulting shape is... The feature map is expanded in the spatial dimension to obtain a shape of... Feature map; Finally, a kernel normalizer is used to make the sum of the weights of the convolution kernel equal to 1. In the content-aware reconstruction module, K, centered at each location, is used in the input feature map. up ×K up The region is convolved with the predicted upsampling kernel to obtain the output value.

4. The content-aware lightweight target detection method according to claim 3, characterized in that: In step S5, the YOLOv5s network has several prediction heads set up in the prediction head network, which are used to perform target detection on features of different scales. 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.

5. A content-aware, lightweight target detection method according to claim 4, characterized in that: In step S6, the YOLOv5s network is trained using the partitioned dataset. After training, the YOLOv5s network is lightweighted using a channel pruning algorithm. The channel pruning algorithm uses L1 regularization of the scaling factor γ of the BN layer to evaluate the importance of the input channels, and includes the following process: First, sparse training is performed on the BN layer to make the BN layer sparse; Then, determine the weight threshold based on the number of BN layers to retain, and prune convolutional channels whose absolute scaling factor is less than the weight threshold. Finally, the pruned model is fine-tuned and trained, and the optimal pruned model is obtained through continuous iteration.

Images