Fine-grained image classification method based on segmentation mask and self-attention neural network

Abstract

The present invention proposes a fine-grained image classification method based on segmentation mask and self-attention neural network, which mainly solves the problems of cumbersome training steps, high training difficulty and inability to accurately locate the position of classified objects when the background noise of the image is too large in the existing fine-grained image classification method. The scheme is as follows: download the divided training and test sample sets from the public fine-grained image data set to obtain the category labels corresponding to each image; construct a fine-grained classification model composed of a loss network, a segmentation mask generation network, and a self-attention neural network cascade; use the training sample set to train the classification model using the gradient descent method; input the test sample set into the trained fine-grained classification model to obtain the classification result of the fine-grained image. The present invention can automatically generate a mask to enhance the image, locate the foreground position of the image, weaken the background noise, and has a simple training strategy, which further improves the classification accuracy. It can be used for intelligent security and unmanned retail commercial activities.

Description

Technical Field

[0001] The present invention belongs to the field of computer vision technology, and in particular relates to a fine-grained image classification method, which can be used for commercial activities such as intelligent security and unmanned retail. Background Art

[0002] With the continuous development of deep learning technology in recent years, more and more mature convolutional neural network models have been applied to the field of computer vision, such as unmanned driving, face recognition, medical image analysis, target tracking, etc. Image classification is an important research topic in the field of computer vision. Traditional image classification mainly distinguishes the large categories of objects contained in the image, such as cats and dogs. These labels are always coarse-grained, while the fine-grained image classification problem is to identify the small categories under the large categories, and its classification labels are more refined. Due to the small classification granularity, images of the same large category are more challenging than traditional image classification problems. For example, distinguishing the types of birds, the styles of vehicles, the varieties of flowers, etc. Since the image sub-classification task has the characteristics of small differences between sub-categories and large differences within sub-categories, the core of fine-grained image classification is to extract the detailed features of the image. For example, birds of different sub-categories have similar features in the same static posture, and the extracted features of birds of the same sub-category are significantly different under different postures, lighting, occlusion, etc., so it is necessary to focus on extracting the features of bird beaks, bird heads, bird wings, etc.

[0003] Fine-grained image classification has a wide range of research needs and application scenarios in both academia and industry. Related research fields mainly involve cars, flowers, birds, airplanes and other fields. In life, there are also a wide range of application scenarios. For example: in security scenarios: road monitoring needs to identify the model and year of the moving vehicle; in emerging unmanned retail scenarios: it is necessary to accurately identify the goods to meet the shopping needs of consumers. Therefore, how to design an accurate and efficient fine-grained image classification algorithm model is of great significance.

[0004] Traditional computer vision uses a strong supervision method of manual annotation to extract features from fine-grained images. In addition to using the image classification label, it also requires the use of annotation boxes to complete the detection of foreground objects and local area positioning for local feature extraction. Although this method can effectively improve classification accuracy, it will consume a lot of manual annotation costs, and the location of manual annotation is not necessarily the area with the most significant detail features. The classification accuracy is greatly affected by the manual annotator, and it is not possible to quickly and effectively identify data sets with a large number of image samples and a large number of image categories, so it is not very practical.

[0005] In recent years, with the development and application of deep learning, fine-grained recognition algorithm models based on weak supervision have gradually become mainstream, mainly using convolutional neural networks for feature extraction. Algorithm models based on deep learning can be roughly divided into two categories: local positioning methods and feature encoding methods. The local positioning method mainly locates the image through the attention mechanism module, that is, the detail area with discriminant features is used as local information, and then merged with the overall image information to output a more accurate classification result. For example, a fine-grained image classification method based on feature fusion disclosed in patent application CN202110179265.2; and a fine-grained image classification method based on a multi-layer focused attention network disclosed in patent application CN202011588241.4 are both methods based on local positioning. Although this method can automatically locate the salient areas of the image and fuse the overall features and detail features of the image, it has the disadvantages of high training difficulty, high training complexity, multiple feature extraction stages, and cumbersome processes; the algorithm model based on feature encoding calculates the outer product of different spatial positions and calculates average pooling of features at different spatial positions to obtain the bilinear features of the image. For example, the high-order feature encoding bilinear feature network (Bilinear-CNN, B-CNN) proposed by Lin et al. has high classification accuracy, but cannot extract the nonlinear relationship between feature channels.

[0006] In the past year, the development of self-attention neural network transformer in the field of vision has made a major breakthrough. The original transformer is the most effective model for natural language processing, mainly used in machine translation, sentiment analysis, information extraction and other tasks. Dosovitskiy et al. proposed the visual self-attention neural network (Vision Transformer, ViT), which was the first to apply the self-attention neural network to the field of computer vision: the network divides the image into fixed-size image slices and then extracts features through the self-attention neural network. Although the transformer has strong feature extraction capabilities, the image division form is too simple. When the background ratio in the image is too large, the target object cannot be effectively located, and the feature extraction effect is poor. Summary of the invention

[0007] The purpose of the present invention is to address the deficiencies of the above-mentioned prior art and propose a fine-grained image classification method based on segmentation mask and self-attention neural network to simplify the network training process. Under the premise of ensuring the accuracy of fine-grained image classification, it can effectively locate the target objects with a large background proportion and noisy background noise and improve the feature extraction effect.

[0008] The technical idea for achieving the purpose of the present invention is: extracting the contour features and approximate position information of the original image through a segmentation network, generating a segmentation mask, and effectively positioning the target object when the background accounts for a large proportion and the background noise is noisy; by fusing the segmentation mask with the original image point-to-point, and inputting the fused image into a self-attention neural network, the overall and local detail features of the image are extracted; by directly outputting the classification result through the last layer classification head of the self-attention neural network, the network training process is simplified.

[0009] According to the above ideas, the implementation scheme of the present invention includes the following:

[0010] (1) Download the divided training sample set and test sample set from the public fine-grained image public dataset and obtain the category label corresponding to each image;

[0011] (2) Building a fine-grained classification model:

[0012] (2a) Establish a loss network consisting of an input layer, a regularization layer, and a center vector update layer in cascade order;

[0013] (2b) Select the segmentation mask generation network and the self-attention neural network, and cascade them with the loss network to form a fine-grained classification model;

[0014] (2c) The total loss function L of the fine-grained classification model is set to be the sum of the cross entropy loss function L1, the center vector loss function L2, and the contrast loss function L3, expressed as follows:

[0015] L=αL1(y,y')+βL2+γL3

[0016] Among them, α, β, and γ are the weights of L1, L2, and L3 respectively, and y and y' represent the real image label and the network predicted image label respectively;

[0017] (3) Using the training sample set, the classification model is trained using the gradient descent method. A trained fine-grained classification model is obtained;

[0018] (4) Input the test sample set into the trained fine-grained classification model to obtain the classification results of the fine-grained image.

[0019] Compared with the prior art, the present invention has the following advantages:

[0020] 1. The fine-grained classification model constructed by the present invention can automatically generate masks for enhanced images, locate the foreground of the image, weaken the background noise, and enable the network to find more discriminative overall and detailed features.

[0021] 2. In the fine-grained image classification described in the present invention, there is no need to perform multi-stage complex training on the classification model. Instead, the fine-grained classification model is directly optimized through the defined total loss. The overall process achieves end-to-end optimization, the training strategy is simple, and the classification accuracy is further improved. BRIEF DESCRIPTION OF THE DRAWINGS

[0022] Figure 1 It is a flow chart for realizing the present invention;

[0023] Figure 2 Schematic diagram of the structure of the fine-grained image classification model constructed in the present invention. DETAILED DESCRIPTION

[0024] The present invention is further described in detail below in conjunction with the accompanying drawings and specific examples.

[0025] Reference Figure 1 , the implementation steps of the present invention are as follows:

[0026] Step 1: Download the divided training sample set and test sample set from the public fine-grained image public dataset and obtain the category label corresponding to each image.

[0027] The datasets used in this example are, but not limited to, public datasets commonly used in the field of sub-classification, including the bird dataset Caltech UCSD Bird-2011 proposed by the California Institute of Technology in 2010, the vehicle dataset Stanford Cars proposed by the Stanford University Artificial Intelligence Laboratory, and the canine dataset Stanford Dogs. These datasets have been divided into training sample sets and test sample sets, and the category labels of the samples are provided, among which:

[0028] The CUB-200-2011 dataset contains 11,788 bird images, including 200 bird subcategories. The training dataset has 5,994 images and the test set has 5,794 images. Each image provides image category information, the bounding box of the bird in the image, the key parts of the bird, and the attribute information of the bird. The experiment in this example only requires the category information of the bird, and no other supervision information is required;

[0029] The Stanford Cars dataset contains 16,185 vehicle images, including 196 types of cars. The training dataset has 8,144 images and the test dataset has 8,041 images. The number of images in each category is similar, and the number of images in each category is consistent. They are mainly classified by car brand, model, and year.

[0030] The Stanford Dogs dataset contains 20,580 dog images, including 120 different dog breeds from all over the world. The training set contains 12,000 images and the test set contains 8,580 images. Each image provides class labels and bounding box information. Similarly, this example only requires the class information of the image.

[0031] Step 2: Build a fine-grained classification model.

[0032] Reference Figure 2 , the specific implementation of this step is as follows:

[0033] (2.1) Establish a loss network consisting of an input layer, a regularization layer, and a center vector update layer in cascade:

[0034] Input layer, used to input the central feature vector c i , when the training dataset is CUB-200-2011, the dimension of the center vector is 200×9216, and the value is all 0. The rows of this matrix represent the number of categories of the fine-grained image public dataset, and the columns are the product of the number of self-attention neural network heads and the number of hidden layers. The dimension is adjusted with different datasets or different model sizes;

[0035] Normalization layer, whose normalization method is 2 norm, is used to update the central feature vector to satisfy the normal distribution;

[0036] The center vector update layer is used to update the center feature vector, and its update weight is θ = 0.05;

[0037] (2.2) A segmentation mask generation network is selected to output a mask with the same size as the original image. It includes a convolutional layer, 6 hierarchical depth aggregation layers, a node fusion layer, and a deconvolution layer, where:

[0038] The six hierarchical depth aggregation layers are used to downsample the image. Each hierarchical depth aggregation layer is a tree structure with a similar structure. The tree depths are 1, 1, 1, 2, 2, 1, respectively. Each tree structure contains a root node and 2-8 basic blocks. Each basic block contains two convolution layers, one regularization layer, and one ReLU activation layer, corresponding to Figure 2 In the first column of the left half, the numbers in the boxes represent the downsampling multiples. The image passes through 6 layers of deep aggregation and outputs 4x, 8x, 16x, and 32x downsampled features respectively.

[0039] The node fusion layer is used to upsample and fuse the nodes output by the layered depth aggregation layer, which includes a depth aggregation upsampling operation and an iterative depth aggregation upsampling operation. Both operations use two layers of deconvolution to complete the upsampling operation of the image, such as Figure 2The dashed arrow in the middle. By upsampling features of different sizes at different stages and fusing them, the representation ability of the model can be enhanced.

[0040] The deconvolution layer is used to upsample the output of the node fusion layer. The final output feature map has a resolution of 448x448. After the sigmoid activation operation, the feature map outputs a mask of the same size as the original image.

[0041] The output dimension sizes of the above layers are shown in Table 1:

[0042] Table 1 Output dimensions of each network layer

[0043]

[0044] 2.3) A self-attention neural network is selected to extract the overall and local features of the image. It includes an image slice embedding layer, a position encoding layer, an L-layer self-attention feedforward encoding layer, and two fully connected layers, where:

[0045] The image slice embedding layer is used to divide the mask-enhanced image x into image slices of size D×D. Where H, W, and C are the length, width, and number of channels of the image, respectively. Therefore, the number of image slices N divided is N = HW / D 2 ;

[0046] The position encoding layer is used to encode the position information of the image. This example experiment uses a standard learnable one-dimensional encoding.

[0047] The L-layer self-attention feedforward encoding layer is used to extract the features of the image. Each layer of the self-attention feedforward encoding layer includes a multi-head attention mechanism module MSA, a feedforward connection module MLP and two layers of normalization operations, wherein the number of heads of the multi-head attention mechanism module is K, and the feedforward connection module includes two layers of fully connected layers.

[0048] The output Z' of the multi-head attention module in each layer of the self-attention feedforward network p , the output Z of the feedforward connection module p , respectively expressed as follows:

[0049] Z' p =MSA(LN(Z p-1 ))+Z p-1 ,

[0050] Z p =MLP(LN(Z' p ))+Z' p ,

[0051] Among them, LN represents the layer normalization operation, Z p-1represents the output of the p-1th self-attention feedforward encoding layer, p = 1...L-1

[0052] L is the total number of self-attention feed-forward encoding layers.

[0053] Since the input of the Lth self-attention feedforward encoding layer is the most discriminative K image patches in the L-1th self-attention feedforward encoding layer, these image patches can be calculated from the attention map generated by the self-attention feedforward encoding layer, which is specifically implemented as follows:

[0054] First, the output of the L-1 self-attention feedforward encoding layer is expressed as:

[0055]

[0056] Among them, Z L-1 is the output of the L-1th self-attention feedforward encoding layer, Hook the L-1th layer classification head, N is the number of divided image slices, represents the Uth image slice output by the L-1th self-attention feedforward encoding layer, where U ranges from 1 to N;

[0057] Similarly, the attention graph list a of the K heads output in the multi-head attention mechanism module in the p-th self-attention feedforward encoding layer is p It is expressed as follows:

[0058]

[0059] in, Represents the attention map output by the Wth head in the multi-head attention mechanism module in the pth self-attention feedforward encoding layer, where W ranges from 1 to K;

[0060] Next, the weights of all heads in the first L-1 layer are fused to obtain the final weight a final ;

[0061]

[0062] Finally, in each header, select only the final The image slice corresponding to the maximum position of the response value in the Lth layer self-attention feedforward encoding layer is finally obtained. local :

[0063]

[0064] Among them A W represents the Wth head of the multi-head self-attention mechanism module, Represents the most discriminative image patch selected by the Wth head of the multi-head attention mechanism module in the L-1th self-attention feedforward encoding layer. The value range of W is 1 to K. In this example, K=12, L=12, and D=14.

[0065] The fully connected layer is used to obtain the final classification prediction result y'.

[0066] 2.4) The segmentation mask generation network, the self-attention neural network, and the loss network are sequentially connected to form a fine-grained classification model, such as Figure 2 shown.

[0067] 2.5) Set the total loss function L of the fine-grained classification model to be the sum of the cross entropy loss function L1, the center vector loss function L2, and the contrast loss function L3, expressed as follows:

[0068] L=αL1(y,y')+βL2+γL3

[0069] Among them, α, β, and γ are the weights of L1, L2, and L3 respectively, and y and y' represent the real image label and the network predicted image label respectively.

[0070] The cross entropy loss function L1, the center vector loss function L2, and the contrast loss function L3 are expressed as follows:

[0071] L1=-[ylogy'+(1-y)log(1-y')]

[0072]

[0073]

[0074] Where N is the batch size, z i With z j Represent the i-th category and the j-th category, α is the loss boundary, y i With y j represents the category of the i-th image and the j-th image, sim(z i , z j ) represents the i-th category image z i and the jth category image z j The cosine similarity between ; Represents y i The characteristic center vector of the class, The vector representing the output of the 0th image patch of the Lth layer of the self-attention neural network.

[0075] Step 3: Use the training sample set to train the classification model using the gradient descent method. The trained fine-grained classification model is obtained:

[0076] (3a) Initialize the parameters of the segmentation mask generation network, self-attention neural network, and loss network:

[0077] The weights of the segmentation mask generation network pre-trained on the ImageNet image dataset are used as its initial weights;

[0078] The weights of the self-attention neural network pre-trained on the ImageNet image dataset are used as its initial weights;

[0079] The loss network is initialized randomly.

[0080] Set the maximum number of iterations of the fine-grained classification model to 10,000, the initial learning rate to 0.03, the learning rate decay strategy to fixed step decay, decaying by 0.1 every 3,000 steps, the optimizer to SGD, and the momentum to 0.9;

[0081] (3b) The images of the training sample set are sequentially cropped, randomly horizontally flipped, randomly erased, and normalized before being sent to the fine-grained classification model;

[0082] (3c) Assume that the weights α and γ of the cross entropy and contrast loss functions are both 1, and the weight β of the center vector loss function is 0.1; calculate the cross entropy loss function L1 based on the prediction result y′ output by the fine-grained classification model and the true label y; and calculate the output category y in the same batch based on the output category y i and the category center vector Calculate the center vector loss function L2; calculate the contrast loss function L3 based on the cosine similarity between image slices of different categories and between image slices of the same category;

[0083] (3d) Substitute the calculated result of (3c) into the total loss function L of the fine-grained classification model, and use the gradient descent algorithm to iteratively update the parameters of the fine-grained classification model;

[0084] (3e) Record the accuracy and loss value of each training iteration, perform validation every 100 iterations, and save the validation model parameters with the best effect until the maximum iteration is 10,000, and then obtain the trained fine-grained classification model.

[0085] Step 4: Input the test sample set into the trained fine-grained classification model to obtain the classification result of the fine-grained image.

[0086] The effect of the present invention can be further illustrated by the following comparative data:

[0087] 5794 test samples of the CUB-200-2011 dataset are selected to verify and compare the effectiveness of the present invention and the existing seven methods, namely DBTNet, S3N, FDL, API-Net, StackedLSTM, ViT, and TransFG. The experimental results are shown in Table 2.

[0088] Table 2 Classification results of different methods on the CUB-200-2011 dataset

[0089] method Recognition accuracy DBTNet 88.1％ S3N 88.5％ FDL 89.1％ API-Net 90.0％ StackedLSTM 90.4％ ViT 90.3％ TransFG 91.4％ The present invention 91.7％

[0090] It can be seen from Table 2 that the recognition accuracy of the present invention is significantly higher than that of the previous mainstream methods, and it also has obvious advantages over other methods.

[0091] The above description is only a specific example of the present invention and does not constitute any limitation to the present invention. Obviously, for professionals in this field, after understanding the content and principles of the present invention, it is possible to make various modifications and changes in form and details without departing from the principles and structures of the present invention, such as modifying or replacing the segmentation mask generation network or the self-attention neural network in the present invention. However, these modifications and changes based on the ideas of the present invention are still within the scope of protection of the claims of the present invention.

Claims

1. A fine-grained image classification method based on segmentation mask and self-attention neural network, characterized in that it includes: (1) Download the divided training sample set and test sample set from the public fine-grained image dataset and obtain the category label corresponding to each image; (2) Building a fine-grained classification model: (2a) Establish a loss network consisting of an input layer, a regularization layer, and a center vector update layer in cascade order; (2b) Select the segmentation mask generation network and the self-attention neural network, and cascade them with the loss network to form a fine-grained classification model; (2c) The total loss function L of the fine-grained classification model is set to be the sum of the cross entropy loss function L1, the center vector loss function L2, and the contrast loss function L3, expressed as follows: L=αL1(y,y')+βL2+γL3 Among them, α, β, γ are the weights of L1, L2, and L3 respectively, y and y' represent the real image label and the network predicted image label respectively; the cross entropy loss function L1, the center vector loss function L2, and the contrast loss function L3 are expressed as follows: L1=-[ylogy'+(1-y)log(1-y')] Where y and y' represent the real image label and the network predicted image label, N is the batch size, and z i With z j Represent the i-th category image and the j-th category image respectively, α′ is the loss boundary, y i With y j represents the category of the i-th image and the j-th image, sim(z i ,z j ) represents the i-th category image z i and the jth category image z j The cosine similarity between ; Represents y i The characteristic center vector of the class, The vector representing the output of the 0th image patch of the Lth layer of the self-attention neural network; (3) Using the training sample set, the classification model is trained using the gradient descent method to obtain a trained fine-grained classification model; (4) Input the test sample set into the trained fine-grained classification model to obtain the classification results of the fine-grained image.

2. The method according to claim 1, wherein: (2a) The loss network established has the following parameters for each layer: The input layer is used to input the central feature vector. Its dimension is 200×9216, and it is a learnable matrix with all values ​​​​set to 0. The rows of this matrix represent the number of categories in the fine-grained image public dataset, and the columns are the product of the number of self-attention neural network heads and the number of hidden layers. The dimension is adjusted with different datasets or different model sizes; Normalization layer, whose normalization method is 2 norm, is used to update the central feature vector to satisfy the normal distribution; The center vector update layer is used to update the center feature vector, and its update weight is β′=0.

05.

3. The method according to claim 1, wherein: The segmentation mask generation network selected in (2b) is used to output a mask with the same size as the original image. It includes a convolutional layer, 6 hierarchical depth aggregation layers, a node fusion layer, and a deconvolution layer, where: The six hierarchical depth aggregation layers are used for downsampling nodes. Each hierarchical depth aggregation layer is a tree structure with a similar structure. The tree depths are 1, 1, 1, 2, 2, 1 respectively. Each tree structure includes a root node and 2-8 basic blocks. Each basic block includes two convolution layers, one regularization layer, and one ReLU activation layer. The node fusion layer is used to upsample and fuse nodes, which includes a depth aggregation upsampling operation and an iterative depth aggregation upsampling operation; The depth aggregation upsampling operation is to fuse the output of the third depth aggregation layer and the output of the fourth depth aggregation layer to generate a first node O1, fuse the output of the fourth depth aggregation layer and the output of the fifth depth aggregation layer to generate a second node O2, fuse the output of the fifth depth aggregation layer and the output of the sixth depth aggregation layer to generate a third node O3, fuse the first node O1 and the second node O2 to generate a fourth node O4, fuse the second node O2 and the third node O3 to generate a fifth node O5, and fuse the fourth node O4 and the fifth node O5 to generate a sixth node O6; The iterative deep aggregation upsampling operation fuses the fifth node O5 and the sixth node O6 to generate the seventh node O7, fuses the seventh node O7 and the third node O3 to generate the eighth node O8, and fuses the eighth node O8 and the output of the third deep aggregation layer to generate the ninth node O9.

4. The method according to claim 1, wherein: The self-attention neural network in (2b) is composed of an image slice embedding layer, a position encoding layer, 12 self-attention feedforward encoding layers, and two fully connected layers connected in sequence. Each self-attention feedforward encoding layer contains a multi-head attention mechanism module MSA and a feedforward connection module MLP. The number of heads of the multi-head attention mechanism module is 12, and the dimension of the feedforward connection module is 3072.

5. The method according to claim 1, wherein: In (3), the fine-grained classification model is trained using the training sample set and the gradient descent method; the implementation is as follows: (3a) Initialize the parameters of the segmentation mask generation network, self-attention neural network, and loss network: The weights of the segmentation mask generation network pre-trained on the ImageNet image dataset are used as its initial weights; The weights of the self-attention neural network pre-trained on the ImageNet image dataset are used as its initial weights; Initialize the loss network randomly; Set the maximum number of iterations of the fine-grained classification model to 10,000, the initial learning rate to 0.03, the learning rate decay strategy to fixed step decay, decaying by 0.1 every 3,000 steps, the optimizer to SGD, and the momentum to 0.9; (3b) The images of the training sample set are sequentially cropped, randomly horizontally flipped, randomly erased, and normalized before being sent to the fine-grained classification model; (3c) Assume that the weights α and γ of the cross entropy and contrast loss functions are both 1, and the weight β of the center vector loss function is 0.1; calculate the cross entropy loss function L1 based on the prediction result y' output by the fine-grained classification model and the true label y; and calculate the cross entropy loss function L2 based on the output category y in the same batch. i and the category center vector c yi Calculate the center vector loss function L2; calculate the contrast loss function L3 based on the cosine similarity between image slices of different categories and between image slices of the same category; (3d) Substitute the calculated result of (3c) into the total loss function L of the fine-grained classification model, and use the gradient descent algorithm to iteratively update the parameters of the fine-grained classification model; (3e) Record the accuracy and loss value of each training iteration, perform verification every 100 iterations, and save the verification model parameters with the best effect; until the maximum iteration is 10,000, a trained fine-grained classification model is obtained.

Images